from __future__ import annotations

from collections import defaultdict, deque
from typing import Any, Union, get_args, get_origin, get_type_hints


class AutoWire:
    """Infers execution edges from task type annotations."""

    def infer_edges(self, task_fns: list) -> list[tuple[str, str]]:
        """Infer task dependencies where return type matches downstream parameter type."""
        edges: set[tuple[str, str]] = set()
        hints_map = {fn.__name__: get_type_hints(fn) for fn in task_fns}

        for producer in task_fns:
            p_hints = hints_map[producer.__name__]
            p_ret = self._unwrap_optional(p_hints.get("return", Any))
            if p_ret is Any:
                continue
            for consumer in task_fns:
                if producer is consumer:
                    continue
                c_hints = hints_map[consumer.__name__]
                for pname, ann in c_hints.items():
                    if pname == "return":
                        continue
                    if self._unwrap_optional(ann) == p_ret:
                        edges.add((producer.__name__, consumer.__name__))
                        break
        return sorted(edges)

    def _unwrap_optional(self, annotation: type) -> type:
        """Unwrap Optional[T] to T when appropriate."""
        origin = get_origin(annotation)
        if origin is Union or str(origin) == "types.UnionType":
            args = [a for a in get_args(annotation) if a is not type(None)]
            if len(args) == 1:
                return args[0]
        return annotation

    def build_dependency_order(self, task_fns: list, explicit_depends_on: dict[str, list[str]]) -> list[str]:
        """Merge inferred and explicit dependencies and return topological order."""
        names = [fn.__name__ for fn in task_fns]
        adjacency: dict[str, list[str]] = defaultdict(list)
        indeg: dict[str, int] = {name: 0 for name in names}

        for a, b in self.infer_edges(task_fns):
            adjacency[a].append(b)
            indeg[b] += 1

        for tgt, deps in explicit_depends_on.items():
            for dep in deps:
                if dep in indeg and tgt in indeg:
                    adjacency[dep].append(tgt)
                    indeg[tgt] += 1

        q = deque([name for name, deg in indeg.items() if deg == 0])
        out: list[str] = []
        while q:
            cur = q.popleft()
            out.append(cur)
            for nxt in adjacency[cur]:
                indeg[nxt] -= 1
                if indeg[nxt] == 0:
                    q.append(nxt)

        if len(out) != len(names):
            unresolved = [name for name in names if name not in out]
            out.extend(unresolved)
        return out
