#
#   Imandra Inc.
#
#   speclogician/modeling/reachability.py
#

from __future__ import annotations

from typing import Any, Dict, List, Tuple, TypeAlias
from pydantic import BaseModel, Field

# A hashable, deterministic edge identifier (ignore witness for diff purposes)
ReachabilityEdgeKey: TypeAlias = Tuple[str, str]


class ScenarioReachabilityEdge(BaseModel):
    src: str
    dst: str
    witness: str | None = None

    def key(self) -> ReachabilityEdgeKey:
        return (self.src, self.dst)


class ScenarioReachabilityGraphDiff(BaseModel):
    """
    Structural diff for reachability graphs.
    Deterministic + lightweight:
      - ignores witness blobs for diff purposes
      - diffs nodes, edge keys, initials, reachable/unreachable sets
    """

    present_before: bool = False
    present_after: bool = False

    nodes_added: list[str] = Field(default_factory=list)
    nodes_removed: list[str] = Field(default_factory=list)

    edges_added: list[ReachabilityEdgeKey] = Field(default_factory=list)
    edges_removed: list[ReachabilityEdgeKey] = Field(default_factory=list)

    initials_added: list[str] = Field(default_factory=list)
    initials_removed: list[str] = Field(default_factory=list)

    reachable_added: list[str] = Field(default_factory=list)
    reachable_removed: list[str] = Field(default_factory=list)

    unreachable_added: list[str] = Field(default_factory=list)
    unreachable_removed: list[str] = Field(default_factory=list)

    def has_changes(self) -> bool:
        return any(
            [
                self.present_before != self.present_after,
                self.nodes_added,
                self.nodes_removed,
                self.edges_added,
                self.edges_removed,
                self.initials_added,
                self.initials_removed,
                self.reachable_added,
                self.reachable_removed,
                self.unreachable_added,
                self.unreachable_removed,
            ]
        )


class ScenarioReachabilityGraph(BaseModel):
    """
    Pure semantic reachability graph between scenarios.
    """

    nodes: list[str] = Field(default_factory=list)
    edges: list[ScenarioReachabilityEdge] = Field(default_factory=list)

    # adjacency
    succ: Dict[str, List[str]] = Field(default_factory=dict)
    pred: Dict[str, List[str]] = Field(default_factory=dict)

    # optional initial analysis
    initials: list[str] = Field(default_factory=list)
    reachable_from_initials: list[str] = Field(default_factory=list)
    unreachable: list[str] = Field(default_factory=list)

    # ----------------------------
    # derived helpers
    # ----------------------------

    def successors(self, name: str) -> list[str]:
        return self.succ.get(name, [])

    def predecessors(self, name: str) -> list[str]:
        return self.pred.get(name, [])

    def is_unreachable(self, name: str) -> bool:
        return name in self.unreachable

    def edge(self, src: str, dst: str) -> ScenarioReachabilityEdge | None:
        for e in self.edges:
            if e.src == src and e.dst == dst:
                return e
        return None

    def explain_edge(
        self,
        src: str,
        dst: str,
        *,
        include_witness_src: bool = True,
        semantics: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        def _fp(s: str) -> str:
            import hashlib

            return hashlib.sha1(s.encode("utf-8")).hexdigest()[:10]

        e = self.edge(src, dst)
        if e is None:
            return {
                "reachable": False,
                "witness_fp": None,
                "witness_src": None,
                "semantics": semantics,
                "notes": "no edge in reachability graph",
            }

        w = e.witness
        return {
            "reachable": True,
            "witness_fp": _fp(w) if w else None,
            "witness_src": (w if (include_witness_src and w) else None),
            "semantics": semantics,
        }

    def shortest_path(self, frm: str, to: str) -> list[str] | None:
        if frm == to:
            return [frm]
        if frm not in self.nodes or to not in self.nodes:
            return None

        from collections import deque

        q: deque[str] = deque([frm])
        parent: dict[str, str | None] = {frm: None}

        while q:
            cur = q.popleft()
            for nxt in self.successors(cur):
                if nxt in parent:
                    continue
                parent[nxt] = cur
                if nxt == to:
                    path = [to]
                    p: str | None = cur
                    while p is not None:
                        path.append(p)
                        p = parent[p]
                    path.reverse()
                    return path
                q.append(nxt)
        return None

    def frontier_gaps(self) -> list[tuple[str, str]]:
        universe = list(self.reachable_from_initials) if self.initials else list(self.nodes)
        out: list[tuple[str, str]] = []
        for src in universe:
            if self.successors(src):
                continue
            if self.initials and src in self.initials:
                out.append((src, "initial_scenario_has_no_successors"))
            else:
                out.append((src, "no_successors_in_graph"))
        return out


def build_reachability_graph(
    *,
    scenario_names: list[str],
    edges: list[ScenarioReachabilityEdge],
    initials: list[str] | None = None,
) -> ScenarioReachabilityGraph:
    names = sorted(set(scenario_names))

    succ: dict[str, list[str]] = {n: [] for n in names}
    pred: dict[str, list[str]] = {n: [] for n in names}

    for e in edges:
        succ.setdefault(e.src, []).append(e.dst)
        pred.setdefault(e.dst, []).append(e.src)

    # normalize deterministically
    for n in sorted(set(list(succ.keys()) + list(pred.keys()) + names)):
        succ.setdefault(n, [])
        pred.setdefault(n, [])
        succ[n] = sorted(set(succ[n]))
        pred[n] = sorted(set(pred[n]))

    names = sorted(set(list(succ.keys()) + list(pred.keys())))
    initials = [n for n in (initials or []) if n in names]

    reachable: set[str] = set()
    queue: list[str] = []

    for n in initials:
        reachable.add(n)
        queue.append(n)

    while queue:
        cur = queue.pop(0)
        for nxt in succ.get(cur, []):
            if nxt not in reachable:
                reachable.add(nxt)
                queue.append(nxt)

    unreachable = sorted(n for n in names if n not in reachable) if initials else []

    return ScenarioReachabilityGraph(
        nodes=names,
        edges=edges,
        succ=succ,
        pred=pred,
        initials=initials,
        reachable_from_initials=sorted(reachable),
        unreachable=unreachable,
    )


def diff_reachability_graph(
    before: ScenarioReachabilityGraph | None,
    after: ScenarioReachabilityGraph | None,
) -> ScenarioReachabilityGraphDiff:
    """
    Compute a deterministic structural diff between two graphs.

    Notes:
      - Witness strings are ignored for edge diffing (we diff by (src,dst) only).
      - If both graphs are None, returns "no changes" diff.
    """
    present_before = before is not None
    present_after = after is not None

    b_nodes = set((before.nodes if before else []) or [])
    a_nodes = set((after.nodes if after else []) or [])

    def _edge_keys(g: ScenarioReachabilityGraph | None) -> set[ReachabilityEdgeKey]:
        if g is None:
            return set()
        return {(e.src, e.dst) for e in (g.edges or [])}

    b_edges = _edge_keys(before)
    a_edges = _edge_keys(after)

    b_initials = set((before.initials if before else []) or [])
    a_initials = set((after.initials if after else []) or [])

    b_reach = set((before.reachable_from_initials if before else []) or [])
    a_reach = set((after.reachable_from_initials if after else []) or [])

    b_unreach = set((before.unreachable if before else []) or [])
    a_unreach = set((after.unreachable if after else []) or [])

    # deterministic list ordering
    nodes_added = sorted(a_nodes - b_nodes)
    nodes_removed = sorted(b_nodes - a_nodes)

    edges_added = sorted(a_edges - b_edges)
    edges_removed = sorted(b_edges - a_edges)

    initials_added = sorted(a_initials - b_initials)
    initials_removed = sorted(b_initials - a_initials)

    reachable_added = sorted(a_reach - b_reach)
    reachable_removed = sorted(b_reach - a_reach)

    unreachable_added = sorted(a_unreach - b_unreach)
    unreachable_removed = sorted(b_unreach - a_unreach)

    return ScenarioReachabilityGraphDiff(
        present_before=present_before,
        present_after=present_after,
        nodes_added=nodes_added,
        nodes_removed=nodes_removed,
        edges_added=edges_added,
        edges_removed=edges_removed,
        initials_added=initials_added,
        initials_removed=initials_removed,
        reachable_added=reachable_added,
        reachable_removed=reachable_removed,
        unreachable_added=unreachable_added,
        unreachable_removed=unreachable_removed,
    )