"""
TODO:
 - @property dfs, bfs
"""
import functools
import typing as ta
import weakref

from .. import cached
from .. import check
from .. import collections as col
from .. import lang


T = ta.TypeVar('T')
NodeT = ta.TypeVar('NodeT')
NodeWalker = ta.Callable[[NodeT], ta.Iterable[NodeT]]
NodeGenerator = ta.Generator[NodeT, None, None]


class NodeError(ta.Generic[NodeT], Exception):
    def __init__(self, node: NodeT, msg: str, *args, **kwargs) -> None:
        super().__init__(msg, *args, **kwargs)  # noqa
        self._node = node

    @property
    def node(self) -> NodeT:
        return self._node


class DuplicateNodeError(NodeError[NodeT]):
    def __init__(self, node: NodeT, *args, **kwargs) -> None:
        super().__init__(node, f'Duplicate node: {node!r}', *args, **kwargs)


class UnknownNodeError(NodeError[NodeT]):
    def __init__(self, node: NodeT, *args, **kwargs) -> None:
        super().__init__(node, f'Unknown node: {node!r}', *args, **kwargs)


class BasicTreeAnalysis(ta.Generic[NodeT]):

    def __init__(
            self,
            root: NodeT,
            walker: NodeWalker[NodeT],
            *,
            identity: bool = False,
    ) -> None:
        super().__init__()

        self._root = root
        self._walker = walker
        self._identity = identity

        self._set_fac: ta.Callable[..., ta.MutableSet[NodeT]] = col.IdentitySet if identity else set
        self._dict_fac: ta.Callable[..., ta.MutableMapping[NodeT, ta.Any]] = col.IdentityKeyDict if identity else dict
        self._idx_seq_fac: ta.Callable[..., col.IndexedSeq[NodeT]] = functools.partial(col.IndexedSeq, identity=identity)  # type: ignore  # noqa

        def walk(cur: NodeT, parent: NodeT | None) -> None:
            check.not_none(cur)
            if cur in node_set:
                raise DuplicateNodeError(cur)

            nodes.append(cur)
            node_set.add(cur)
            if parent is None:
                check.state(cur is root)
            elif parent not in node_set:
                raise UnknownNodeError(parent)

            parents_by_node[cur] = parent

            children_by_node[cur] = children = list(walker(cur))
            child_sets_by_node[cur] = self._set_fac(children)
            for child in children:
                walk(child, cur)

        nodes: list[NodeT] = []
        node_set: ta.MutableSet[NodeT] = self._set_fac()
        children_by_node: ta.MutableMapping[NodeT | None, ta.Sequence[NodeT]] = self._dict_fac()
        child_sets_by_node: ta.MutableMapping[NodeT | None, ta.AbstractSet[NodeT]] = self._dict_fac()
        parents_by_node: ta.MutableMapping[NodeT, NodeT | None] = self._dict_fac()

        children_by_node[None] = [root]
        child_sets_by_node[None] = self._set_fac([root])

        walk(root, None)

        self._nodes = self._idx_seq_fac(nodes)
        self._node_set: ta.AbstractSet[NodeT] = node_set
        self._children_by_node: ta.Mapping[NodeT | None, col.IndexedSeq[NodeT]] = self._dict_fac(
            [(n, self._idx_seq_fac(cs)) for n, cs in children_by_node.items()])
        self._child_sets_by_node: ta.Mapping[NodeT | None, ta.AbstractSet[NodeT]] = child_sets_by_node
        self._parents_by_node: ta.Mapping[NodeT, NodeT | None] = parents_by_node

    @property
    def root(self) -> NodeT:
        return self._root

    @property
    def nodes(self) -> col.IndexedSeq[NodeT]:
        return self._nodes

    @property
    def walker(self) -> NodeWalker[NodeT]:
        return self._walker

    @property
    def identity(self) -> bool:
        return self._identity

    @property
    def node_set(self) -> ta.AbstractSet[NodeT]:
        return self._node_set

    @property
    def children_by_node(self) -> ta.Mapping[NodeT | None, col.IndexedSeq[NodeT]]:
        return self._children_by_node

    @property
    def child_sets_by_node(self) -> ta.Mapping[NodeT | None, ta.AbstractSet[NodeT]]:
        return self._child_sets_by_node

    @property
    def parents_by_node(self) -> ta.Mapping[NodeT, NodeT | None]:
        return self._parents_by_node

    @classmethod
    def from_parents(
            cls,
            src: ta.Mapping[NodeT, NodeT | None] | ta.Iterable[tuple[NodeT, NodeT | None]],
            *,
            identity: bool = False,
            **kwargs,
    ) -> 'BasicTreeAnalysis[NodeT]':
        pairs: ta.Sequence[tuple[NodeT, NodeT]]
        if isinstance(src, ta.Mapping):
            pairs = list(src.items())  # type: ignore
        elif isinstance(src, ta.Iterable):
            pairs = list(src)  # type: ignore
        else:
            raise TypeError(src)

        pairs = [(check.not_none(n), p) for n, p in pairs]

        root = check.single([n for n, p in pairs if p is None])  # noqa

        children_by_node: ta.MutableMapping[NodeT, ta.MutableSequence[NodeT]] = col.IdentityKeyDict() if identity else {}  # noqa
        for n, _ in pairs:
            children_by_node[n] = []
        for n, p in pairs:
            if p is not None:
                children_by_node[p].append(n)

        return cls(
            root,
            children_by_node.__getitem__,
            identity=identity,
            **kwargs,
        )

    @classmethod
    def from_children(
            cls,
            src: ta.Mapping[NodeT, ta.Iterable[NodeT]] | ta.Iterable[tuple[NodeT, ta.Iterable[NodeT]]],
            *,
            identity: bool = False,
            **kwargs,
    ) -> 'BasicTreeAnalysis[NodeT]':
        pairs: ta.Sequence[tuple[NodeT, ta.Sequence[NodeT]]]
        if isinstance(src, ta.Mapping):
            pairs = list(src.items())  # type: ignore
        elif isinstance(src, ta.Iterable):
            pairs = list(src)  # type: ignore
        else:
            raise TypeError(src)

        pairs = [(check.not_none(n), [check.not_none(c) for c in cs]) for n, cs in pairs]

        children_by_node: ta.MutableMapping[NodeT, ta.Sequence[NodeT]] = col.IdentityKeyDict() if identity else {}
        parents_by_node: ta.MutableMapping[NodeT, NodeT] = col.IdentityKeyDict() if identity else {}
        for n, cs in pairs:
            check.not_in(n, children_by_node)
            children_by_node[n] = cs
            for c in cs:
                check.not_in(c, parents_by_node)
                parents_by_node[c] = n

        e: ta.Any
        d: ta.Any
        if identity:
            e, d = id, col.make_map(((id(n), n) for n, _ in pairs), strict=True)
        else:
            e, d = lang.identity, lang.identity
        tsd = {e(n): {e(p)} for n, p in parents_by_node.items()}
        ts = list(col.mut_toposort(tsd))
        root = d(check.single(ts[0]))

        return cls(
            root,
            children_by_node.__getitem__,
            identity=identity,
            **kwargs,
        )

    @cached.property
    def _node_sets_by_type(self) -> ta.MutableMapping[type, ta.AbstractSet[NodeT]]:
        return weakref.WeakKeyDictionary()

    def get_node_type_set(self, ty: type[T]) -> ta.AbstractSet[T]:
        try:
            return self._node_sets_by_type[ty]  # type: ignore
        except KeyError:
            ret = self._node_sets_by_type[ty] = self._set_fac(n for n in self.nodes if isinstance(n, ty))
            return ret  # type: ignore

    def iter_ancestors(self, node: NodeT) -> NodeGenerator[NodeT]:
        cur: NodeT | None = node
        while True:
            cur = self.parents_by_node.get(cur)  # type: ignore
            if cur is None:
                break
            yield cur

    def get_lineage(self, node: NodeT) -> col.IndexedSeq[NodeT]:
        return self._idx_seq_fac(reversed([node, *self.iter_ancestors(node)]))

    def get_first_parent_of_type(self, node: NodeT, ty: type[T]) -> T | None:
        for cur in self.iter_ancestors(node):
            if isinstance(cur, ty):
                return cur
        return None

    @cached.property
    def depths_by_node(self) -> ta.Mapping[NodeT, int]:
        def rec(n, d):
            dct[n] = d
            for c in self._children_by_node[n]:
                rec(c, d + 1)
        dct: ta.MutableMapping[NodeT, int] = self._dict_fac()
        rec(self._root, 0)
        return dct
