from __future__ import annotations

from collections import deque
from typing import Any

from ..schema import SuggestedFlow
from .graph_builder import build_module_adjacency


def _shortest_path(
    out_adj: dict[str, dict[str, int]],
    start: str,
    targets: set[str],
    *,
    max_len: int = 6,
) -> list[str] | None:
    if start in targets:
        return [start]
    q: deque[str] = deque([start])
    prev: dict[str, str | None] = {start: None}
    depth: dict[str, int] = {start: 0}

    while q:
        cur = q.popleft()
        d = depth.get(cur, 0)
        if d >= max_len:
            continue
        for nxt in out_adj.get(cur, {}).keys():
            if nxt in prev:
                continue
            prev[nxt] = cur
            depth[nxt] = d + 1
            if nxt in targets:
                # reconstruct
                path: list[str] = [nxt]
                p = cur
                while p is not None:
                    path.append(p)
                    p = prev.get(p)
                path.reverse()
                return path
            q.append(nxt)
    return None


def _edges_evidence(path: list[str], out_adj: dict[str, dict[str, int]]) -> list[dict[str, Any]]:
    ev: list[dict[str, Any]] = []
    for a, b in zip(path, path[1:]):
        ev.append({"from_module": a, "to_module": b, "weight": int(out_adj.get(a, {}).get(b, 0))})
    return ev


def suggest_flows(
    *,
    modules: list[dict[str, Any]],
    module_edges: list[dict[str, Any]],
    fan_in_w: dict[str, int],
    fan_out_w: dict[str, int],
    limit: int = 10,
) -> list[SuggestedFlow]:
    out_adj, _ = build_module_adjacency(module_edges)  # type: ignore[arg-type]
    mids = [str(m.get("module_id") or "") for m in modules if str(m.get("module_id") or "")]

    degree_w = {m: int(fan_in_w.get(m, 0)) + int(fan_out_w.get(m, 0)) for m in mids}

    entries = sorted(
        mids,
        key=lambda m: (int(fan_in_w.get(m, 0)), -int(fan_out_w.get(m, 0)), -int(degree_w.get(m, 0)), m),
    )
    entries = [m for m in entries if int(fan_out_w.get(m, 0)) > 0][:8]

    shared = sorted(
        mids,
        key=lambda m: (-int(fan_in_w.get(m, 0)), int(fan_out_w.get(m, 0)), m),
    )
    shared = [m for m in shared if int(fan_in_w.get(m, 0)) >= 3][:8]

    core = sorted(mids, key=lambda m: (-int(degree_w.get(m, 0)), m))[:10]

    flows: list[SuggestedFlow] = []
    seen: set[tuple[str, ...]] = set()

    def add_flow(mods: list[str], name: str) -> None:
        key = tuple(mods)
        if len(mods) < 2 or key in seen:
            return
        seen.add(key)
        flows.append(
            {
                "name": name,
                "modules": mods,
                "evidence_edges": _edges_evidence(mods, out_adj),
            }
        )

    for entry in entries:
        if len(flows) >= max(3, int(limit)):
            break
        p1 = _shortest_path(out_adj, entry, set(core), max_len=5)
        if not p1:
            continue
        mid = p1[-1]
        p2 = _shortest_path(out_adj, mid, set(shared), max_len=4) if shared else None
        if p2 and p2[0] == mid:
            mods = p1 + p2[1:]
        else:
            mods = p1
        add_flow(mods, f"{entry} flow")

    # If still not enough, add core->shared or entry->shared directly
    if len(flows) < 3:
        for c in core:
            if len(flows) >= max(3, int(limit)):
                break
            p = _shortest_path(out_adj, c, set(shared), max_len=5) if shared else None
            if p:
                add_flow(p, f"{c} to shared")

    return flows[: max(0, int(limit))]

