from typing import Any, Callable, Iterable, List, Tuple, TypeVar, Union

from bigtree.node.basenode import BaseNode
from bigtree.node.dagnode import DAGNode
from bigtree.node.node import Node
from bigtree.utils.exceptions import SearchError
from bigtree.utils.iterators import preorder_iter

__all__ = [
    "findall",
    "find",
    "find_name",
    "find_names",
    "find_relative_path",
    "find_full_path",
    "find_path",
    "find_paths",
    "find_attr",
    "find_attrs",
    "find_children",
    "find_child",
    "find_child_by_name",
]


T = TypeVar("T", bound=BaseNode)
NodeT = TypeVar("NodeT", bound=Node)
DAGNodeT = TypeVar("DAGNodeT", bound=DAGNode)


def findall(
    tree: T,
    condition: Callable[[T], bool],
    max_depth: int = 0,
    min_count: int = 0,
    max_count: int = 0,
) -> Tuple[T, ...]:
    """
    Search tree for nodes matching condition (callable function).

    Examples:
        >>> from bigtree import Node, findall
        >>> root = Node("a", age=90)
        >>> b = Node("b", age=65, parent=root)
        >>> c = Node("c", age=60, parent=root)
        >>> d = Node("d", age=40, parent=c)
        >>> findall(root, lambda node: node.age > 62)
        (Node(/a, age=90), Node(/a/b, age=65))

    Args:
        tree (BaseNode): tree to search
        condition (Callable): function that takes in node as argument, returns node if condition evaluates to `True`
        max_depth (int): maximum depth to search for, based on the `depth` attribute, defaults to None
        min_count (int): checks for minimum number of occurrences,
            raise SearchError if the number of results do not meet min_count, defaults to None
        max_count (int): checks for maximum number of occurrences,
            raise SearchError if the number of results do not meet min_count, defaults to None

    Returns:
        (Tuple[BaseNode, ...])
    """
    result = tuple(preorder_iter(tree, filter_condition=condition, max_depth=max_depth))
    if min_count and len(result) < min_count:
        raise SearchError(
            f"Expected more than {min_count} element(s), found {len(result)} elements\n{result}"
        )
    if max_count and len(result) > max_count:
        raise SearchError(
            f"Expected less than {max_count} element(s), found {len(result)} elements\n{result}"
        )
    return result


def find(tree: T, condition: Callable[[T], bool], max_depth: int = 0) -> T:
    """
    Search tree for *single node* matching condition (callable function).

    Examples:
        >>> from bigtree import Node, find
        >>> root = Node("a", age=90)
        >>> b = Node("b", age=65, parent=root)
        >>> c = Node("c", age=60, parent=root)
        >>> d = Node("d", age=40, parent=c)
        >>> find(root, lambda node: node.age == 65)
        Node(/a/b, age=65)
        >>> find(root, lambda node: node.age > 5)
        Traceback (most recent call last):
            ...
        bigtree.utils.exceptions.SearchError: Expected less than 1 element(s), found 4 elements
        (Node(/a, age=90), Node(/a/b, age=65), Node(/a/c, age=60), Node(/a/c/d, age=40))

    Args:
        tree (BaseNode): tree to search
        condition (Callable): function that takes in node as argument, returns node if condition evaluates to `True`
        max_depth (int): maximum depth to search for, based on the `depth` attribute, defaults to None

    Returns:
        (BaseNode)
    """
    result = findall(tree, condition, max_depth, max_count=1)
    if result:
        return result[0]


def find_name(tree: NodeT, name: str, max_depth: int = 0) -> NodeT:
    """
    Search tree for single node matching name attribute.

    Examples:
        >>> from bigtree import Node, find_name
        >>> root = Node("a", age=90)
        >>> b = Node("b", age=65, parent=root)
        >>> c = Node("c", age=60, parent=root)
        >>> d = Node("d", age=40, parent=c)
        >>> find_name(root, "c")
        Node(/a/c, age=60)

    Args:
        tree (Node): tree to search
        name (str): value to match for name attribute
        max_depth (int): maximum depth to search for, based on the `depth` attribute, defaults to None

    Returns:
        (Node)
    """
    return find(tree, lambda node: node.node_name == name, max_depth)


def find_names(tree: NodeT, name: str, max_depth: int = 0) -> Iterable[NodeT]:
    """
    Search tree for multiple node(s) matching name attribute.

    Examples:
        >>> from bigtree import Node, find_names
        >>> root = Node("a", age=90)
        >>> b = Node("b", age=65, parent=root)
        >>> c = Node("c", age=60, parent=root)
        >>> d = Node("b", age=40, parent=c)
        >>> find_names(root, "c")
        (Node(/a/c, age=60),)
        >>> find_names(root, "b")
        (Node(/a/b, age=65), Node(/a/c/b, age=40))

    Args:
        tree (Node): tree to search
        name (str): value to match for name attribute
        max_depth (int): maximum depth to search for, based on the `depth` attribute, defaults to None

    Returns:
        (Iterable[Node])
    """
    return findall(tree, lambda node: node.node_name == name, max_depth)


def find_relative_path(tree: NodeT, path_name: str) -> Iterable[NodeT]:
    r"""
    Search tree for single node matching relative path attribute.

    - Supports unix folder expression for relative path, i.e., '../../node_name'
    - Supports wildcards, i.e., '\*/node_name'
    - If path name starts with leading separator symbol, it will start at root node.

    Examples:
        >>> from bigtree import Node, find_relative_path
        >>> root = Node("a", age=90)
        >>> b = Node("b", age=65, parent=root)
        >>> c = Node("c", age=60, parent=root)
        >>> d = Node("d", age=40, parent=c)
        >>> find_relative_path(d, "..")
        (Node(/a/c, age=60),)
        >>> find_relative_path(d, "../../b")
        (Node(/a/b, age=65),)
        >>> find_relative_path(d, "../../*")
        (Node(/a/b, age=65), Node(/a/c, age=60))

    Args:
        tree (Node): tree to search
        path_name (str): value to match (relative path) of path_name attribute

    Returns:
        (Iterable[Node])
    """
    sep = tree.sep
    if path_name.startswith(sep):
        resolved_node = find_full_path(tree, path_name)
        return (resolved_node,)
    path_name = path_name.rstrip(sep).lstrip(sep)
    path_list = path_name.split(sep)
    wildcard_indicator = "*" in path_name
    resolved_nodes: List[NodeT] = []

    def resolve(node: NodeT, path_idx: int) -> None:
        """Resolve node based on path name

        Args:
            node (Node): current node
            path_idx (int): current index in path_list
        """
        if path_idx == len(path_list):
            resolved_nodes.append(node)
        else:
            path_component = path_list[path_idx]
            if path_component == ".":
                resolve(node, path_idx + 1)
            elif path_component == "..":
                if node.is_root:
                    raise SearchError("Invalid path name. Path goes beyond root node.")
                resolve(node.parent, path_idx + 1)
            elif path_component == "*":
                for child in node.children:
                    resolve(child, path_idx + 1)
            else:
                node = find_child_by_name(node, path_component)
                if not node:
                    if not wildcard_indicator:
                        raise SearchError(
                            f"Invalid path name. Node {path_component} cannot be found."
                        )
                else:
                    resolve(node, path_idx + 1)

    resolve(tree, 0)

    return tuple(resolved_nodes)


def find_full_path(tree: NodeT, path_name: str) -> NodeT:
    """
    Search tree for single node matching path attribute.

    - Path name can be with or without leading tree path separator symbol.
    - Path name must be full path, works similar to `find_path` but faster.

    Examples:
        >>> from bigtree import Node, find_full_path
        >>> root = Node("a", age=90)
        >>> b = Node("b", age=65, parent=root)
        >>> c = Node("c", age=60, parent=root)
        >>> d = Node("d", age=40, parent=c)
        >>> find_full_path(root, "/a/c/d")
        Node(/a/c/d, age=40)

    Args:
        tree (Node): tree to search
        path_name (str): value to match (full path) of path_name attribute

    Returns:
        (Node)
    """
    sep = tree.sep
    path_list = path_name.rstrip(sep).lstrip(sep).split(sep)
    if path_list[0] != tree.root.node_name:
        raise ValueError(
            f"Path {path_name} does not match the root node name {tree.root.node_name}"
        )
    parent_node = tree.root
    child_node = parent_node
    for child_name in path_list[1:]:
        child_node = find_child_by_name(parent_node, child_name)
        if not child_node:
            break
        parent_node = child_node
    return child_node


def find_path(tree: NodeT, path_name: str) -> NodeT:
    """
    Search tree for single node matching path attribute.

    - Path name can be with or without leading tree path separator symbol.
    - Path name can be full path or partial path (trailing part of path) or node name.

    Examples:
        >>> from bigtree import Node, find_path
        >>> root = Node("a", age=90)
        >>> b = Node("b", age=65, parent=root)
        >>> c = Node("c", age=60, parent=root)
        >>> d = Node("d", age=40, parent=c)
        >>> find_path(root, "c")
        Node(/a/c, age=60)
        >>> find_path(root, "/c")
        Node(/a/c, age=60)

    Args:
        tree (Node): tree to search
        path_name (str): value to match (full path) or trailing part (partial path) of path_name attribute

    Returns:
        (Node)
    """
    path_name = path_name.rstrip(tree.sep)
    return find(tree, lambda node: node.path_name.endswith(path_name))


def find_paths(tree: NodeT, path_name: str) -> Tuple[NodeT, ...]:
    """
    Search tree for multiple nodes matching path attribute.

    - Path name can be with or without leading tree path separator symbol.
    - Path name can be partial path (trailing part of path) or node name.

    Examples:
        >>> from bigtree import Node, find_paths
        >>> root = Node("a", age=90)
        >>> b = Node("b", age=65, parent=root)
        >>> c = Node("c", age=60, parent=root)
        >>> d = Node("c", age=40, parent=c)
        >>> find_paths(root, "/a/c")
        (Node(/a/c, age=60),)
        >>> find_paths(root, "/c")
        (Node(/a/c, age=60), Node(/a/c/c, age=40))

    Args:
        tree (Node): tree to search
        path_name (str): value to match (full path) or trailing part (partial path) of path_name attribute

    Returns:
        (Tuple[Node, ...])
    """
    path_name = path_name.rstrip(tree.sep)
    return findall(tree, lambda node: node.path_name.endswith(path_name))


def find_attr(
    tree: BaseNode, attr_name: str, attr_value: Any, max_depth: int = 0
) -> BaseNode:
    """
    Search tree for single node matching custom attribute.

    Examples:
        >>> from bigtree import Node, find_attr
        >>> root = Node("a", age=90)
        >>> b = Node("b", age=65, parent=root)
        >>> c = Node("c", age=60, parent=root)
        >>> d = Node("d", age=40, parent=c)
        >>> find_attr(root, "age", 65)
        Node(/a/b, age=65)

    Args:
        tree (BaseNode): tree to search
        attr_name (str): attribute name to perform matching
        attr_value (Any): value to match for attr_name attribute
        max_depth (int): maximum depth to search for, based on the `depth` attribute, defaults to None

    Returns:
        (BaseNode)
    """
    return find(
        tree,
        lambda node: bool(node.get_attr(attr_name) == attr_value),
        max_depth,
    )


def find_attrs(
    tree: BaseNode, attr_name: str, attr_value: Any, max_depth: int = 0
) -> Tuple[BaseNode, ...]:
    """
    Search tree for node(s) matching custom attribute.

    Examples:
        >>> from bigtree import Node, find_attrs
        >>> root = Node("a", age=90)
        >>> b = Node("b", age=65, parent=root)
        >>> c = Node("c", age=65, parent=root)
        >>> d = Node("d", age=40, parent=c)
        >>> find_attrs(root, "age", 65)
        (Node(/a/b, age=65), Node(/a/c, age=65))

    Args:
        tree (BaseNode): tree to search
        attr_name (str): attribute name to perform matching
        attr_value (Any): value to match for attr_name attribute
        max_depth (int): maximum depth to search for, based on the `depth` attribute, defaults to None

    Returns:
        (Tuple[BaseNode, ...])
    """
    return findall(
        tree,
        lambda node: bool(node.get_attr(attr_name) == attr_value),
        max_depth,
    )


def find_children(
    tree: Union[T, DAGNodeT],
    condition: Callable[[Union[T, DAGNodeT]], bool],
    min_count: int = 0,
    max_count: int = 0,
) -> Tuple[Union[T, DAGNodeT], ...]:
    """
    Search children for nodes matching condition (callable function).

    Examples:
        >>> from bigtree import Node, find_children
        >>> root = Node("a", age=90)
        >>> b = Node("b", age=65, parent=root)
        >>> c = Node("c", age=60, parent=root)
        >>> d = Node("d", age=40, parent=c)
        >>> find_children(root, lambda node: node.age > 30)
        (Node(/a/b, age=65), Node(/a/c, age=60))

    Args:
        tree (BaseNode/DAGNode): tree to search for its children
        condition (Callable): function that takes in node as argument, returns node if condition evaluates to `True`
        min_count (int): checks for minimum number of occurrences,
            raise SearchError if the number of results do not meet min_count, defaults to None
        max_count (int): checks for maximum number of occurrences,
            raise SearchError if the number of results do not meet min_count, defaults to None

    Returns:
        (BaseNode/DAGNode)
    """
    result = tuple([node for node in tree.children if node and condition(node)])
    if min_count and len(result) < min_count:
        raise SearchError(
            f"Expected more than {min_count} element(s), found {len(result)} elements\n{result}"
        )
    if max_count and len(result) > max_count:
        raise SearchError(
            f"Expected less than {max_count} element(s), found {len(result)} elements\n{result}"
        )
    return result


def find_child(
    tree: Union[T, DAGNodeT],
    condition: Callable[[Union[T, DAGNodeT]], bool],
) -> Union[T, DAGNodeT]:
    """
    Search children for *single node* matching condition (callable function).

    Examples:
        >>> from bigtree import Node, find_child
        >>> root = Node("a", age=90)
        >>> b = Node("b", age=65, parent=root)
        >>> c = Node("c", age=60, parent=root)
        >>> d = Node("d", age=40, parent=c)
        >>> find_child(root, lambda node: node.age > 62)
        Node(/a/b, age=65)

    Args:
        tree (BaseNode/DAGNode): tree to search for its child
        condition (Callable): function that takes in node as argument, returns node if condition evaluates to `True`

    Returns:
        (BaseNode/DAGNode)
    """
    result = find_children(tree, condition, max_count=1)
    if result:
        return result[0]


def find_child_by_name(
    tree: Union[NodeT, DAGNodeT], name: str
) -> Union[NodeT, DAGNodeT]:
    """
    Search tree for single node matching name attribute.

    Examples:
        >>> from bigtree import Node, find_child_by_name
        >>> root = Node("a", age=90)
        >>> b = Node("b", age=65, parent=root)
        >>> c = Node("c", age=60, parent=root)
        >>> d = Node("d", age=40, parent=c)
        >>> find_child_by_name(root, "c")
        Node(/a/c, age=60)
        >>> find_child_by_name(c, "d")
        Node(/a/c/d, age=40)

    Args:
        tree (Node/DAGNode): tree to search, parent node
        name (str): value to match for name attribute, child node

    Returns:
        (Node/DAGNode)
    """
    return find_child(tree, lambda node: node.node_name == name)