from __future__ import annotations

import itertools
from collections import namedtuple
from typing import (
    TYPE_CHECKING,
    Callable,
    Generic,
    Iterable,
    Iterator,
    Mapping,
    NamedTuple,
    Sequence,
    TypeVar,
    Union,
)

KT = TypeVar("KT")  # Identifier.
RT = TypeVar("RT")  # Requirement.
CT = TypeVar("CT")  # Candidate.

Matches = Union[Iterable[CT], Callable[[], Iterable[CT]]]

if TYPE_CHECKING:
    from .resolvers.criterion import Criterion

    class RequirementInformation(NamedTuple, Generic[RT, CT]):
        requirement: RT
        parent: CT | None

    class State(NamedTuple, Generic[RT, CT, KT]):
        """Resolution state in a round."""

        mapping: dict[KT, CT]
        criteria: dict[KT, Criterion[RT, CT]]
        backtrack_causes: list[RequirementInformation[RT, CT]]

else:
    RequirementInformation = namedtuple(
        "RequirementInformation", ["requirement", "parent"]
    )
    State = namedtuple("State", ["mapping", "criteria", "backtrack_causes"])


class DirectedGraph(Generic[KT]):
    """A graph structure with directed edges."""

    def __init__(self) -> None:
        self._vertices: set[KT] = set()
        self._forwards: dict[KT, set[KT]] = {}  # <key> -> Set[<key>]
        self._backwards: dict[KT, set[KT]] = {}  # <key> -> Set[<key>]

    def __iter__(self) -> Iterator[KT]:
        return iter(self._vertices)

    def __len__(self) -> int:
        return len(self._vertices)

    def __contains__(self, key: KT) -> bool:
        return key in self._vertices

    def copy(self) -> DirectedGraph[KT]:
        """Return a shallow copy of this graph."""
        other = type(self)()
        other._vertices = set(self._vertices)
        other._forwards = {k: set(v) for k, v in self._forwards.items()}
        other._backwards = {k: set(v) for k, v in self._backwards.items()}
        return other

    def add(self, key: KT) -> None:
        """Add a new vertex to the graph."""
        if key in self._vertices:
            raise ValueError("vertex exists")
        self._vertices.add(key)
        self._forwards[key] = set()
        self._backwards[key] = set()

    def remove(self, key: KT) -> None:
        """Remove a vertex from the graph, disconnecting all edges from/to it."""
        self._vertices.remove(key)
        for f in self._forwards.pop(key):
            self._backwards[f].remove(key)
        for t in self._backwards.pop(key):
            self._forwards[t].remove(key)

    def connected(self, f: KT, t: KT) -> bool:
        return f in self._backwards[t] and t in self._forwards[f]

    def connect(self, f: KT, t: KT) -> None:
        """Connect two existing vertices.

        Nothing happens if the vertices are already connected.
        """
        if t not in self._vertices:
            raise KeyError(t)
        self._forwards[f].add(t)
        self._backwards[t].add(f)

    def iter_edges(self) -> Iterator[tuple[KT, KT]]:
        for f, children in self._forwards.items():
            for t in children:
                yield f, t

    def iter_children(self, key: KT) -> Iterator[KT]:
        return iter(self._forwards[key])

    def iter_parents(self, key: KT) -> Iterator[KT]:
        return iter(self._backwards[key])


class IteratorMapping(Mapping[KT, Iterator[CT]], Generic[RT, CT, KT]):
    def __init__(
        self,
        mapping: Mapping[KT, RT],
        accessor: Callable[[RT], Iterable[CT]],
        appends: Mapping[KT, Iterable[CT]] | None = None,
    ) -> None:
        self._mapping = mapping
        self._accessor = accessor
        self._appends: Mapping[KT, Iterable[CT]] = appends or {}

    def __repr__(self) -> str:
        return "IteratorMapping({!r}, {!r}, {!r})".format(
            self._mapping,
            self._accessor,
            self._appends,
        )

    def __bool__(self) -> bool:
        return bool(self._mapping or self._appends)

    def __contains__(self, key: object) -> bool:
        return key in self._mapping or key in self._appends

    def __getitem__(self, k: KT) -> Iterator[CT]:
        try:
            v = self._mapping[k]
        except KeyError:
            return iter(self._appends[k])
        return itertools.chain(self._accessor(v), self._appends.get(k, ()))

    def __iter__(self) -> Iterator[KT]:
        more = (k for k in self._appends if k not in self._mapping)
        return itertools.chain(self._mapping, more)

    def __len__(self) -> int:
        more = sum(1 for k in self._appends if k not in self._mapping)
        return len(self._mapping) + more


class _FactoryIterableView(Iterable[RT]):
    """Wrap an iterator factory returned by `find_matches()`.

    Calling `iter()` on this class would invoke the underlying iterator
    factory, making it a "collection with ordering" that can be iterated
    through multiple times, but lacks random access methods presented in
    built-in Python sequence types.
    """

    def __init__(self, factory: Callable[[], Iterable[RT]]) -> None:
        self._factory = factory
        self._iterable: Iterable[RT] | None = None

    def __repr__(self) -> str:
        return f"{type(self).__name__}({list(self)})"

    def __bool__(self) -> bool:
        try:
            next(iter(self))
        except StopIteration:
            return False
        return True

    def __iter__(self) -> Iterator[RT]:
        iterable = self._factory() if self._iterable is None else self._iterable
        self._iterable, current = itertools.tee(iterable)
        return current


class _SequenceIterableView(Iterable[RT]):
    """Wrap an iterable returned by find_matches().

    This is essentially just a proxy to the underlying sequence that provides
    the same interface as `_FactoryIterableView`.
    """

    def __init__(self, sequence: Sequence[RT]):
        self._sequence = sequence

    def __repr__(self) -> str:
        return f"{type(self).__name__}({self._sequence})"

    def __bool__(self) -> bool:
        return bool(self._sequence)

    def __iter__(self) -> Iterator[RT]:
        return iter(self._sequence)


def build_iter_view(matches: Matches[CT]) -> Iterable[CT]:
    """Build an iterable view from the value returned by `find_matches()`."""
    if callable(matches):
        return _FactoryIterableView(matches)
    if not isinstance(matches, Sequence):
        matches = list(matches)
    return _SequenceIterableView(matches)


IterableView = Iterable