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

from __future__ import annotations

from collections import deque
from dataclasses import dataclass
from typing import Any, Mapping, Sequence, TypedDict


# NOTE:
# - These helpers are intentionally "thin": they do NOT call ImandraX.
# - They assume spec.scenario_reachability was already computed upstream.
# - They are tolerant to model evolution (missing fields, different names, etc).
# - They prefer using precomputed graph adjacency (succ/pred/unreachable/reachable_from_initials)
#   when available, and fall back to edges-based inference otherwise.


# -----------------------------------------------------------------------------
# Types (lightweight, avoids import cycles)
# -----------------------------------------------------------------------------

class EdgeExplanation(TypedDict, total=False):
    reachable: bool
    witness_fp: str | None
    witness_src: str | None
    semantics: Mapping[str, Any] | None
    notes: str | None


@dataclass(frozen=True)
class FrontierGap:
    src: str
    missing_reason: str


# -----------------------------------------------------------------------------
# Internal helpers
# -----------------------------------------------------------------------------

def _graph(spec: Any) -> Any | None:
    return getattr(spec, "scenario_reachability", None)


def _edges(g: Any) -> list[Any]:
    return list(getattr(g, "edges", []) or [])


def _edge_src(e: Any) -> str:
    return str(getattr(e, "src", getattr(e, "scenario_from", "")) or "")


def _edge_dst(e: Any) -> str:
    return str(getattr(e, "dst", getattr(e, "scenario_to", "")) or "")


def _edge_reachable(e: Any) -> bool:
    v = getattr(e, "reachable", None)
    if v is None:
        v = getattr(e, "is_reachable", None)
    if v is None:
        # In the canonical graph model, presence in edges implies reachable,
        # but keep conservative default:
        v = True
    return bool(v)


def _edge_witness_src(e: Any) -> str | None:
    # canonical field is `witness`; tolerate older names
    w = getattr(e, "witness", None)
    if isinstance(w, str) and w.strip():
        return w
    w = getattr(e, "witness_src", None)
    if isinstance(w, str) and w.strip():
        return w
    w = getattr(e, "model_src", None)
    if isinstance(w, str) and w.strip():
        return w
    return None


def _fp(s: str) -> str:
    import hashlib
    return hashlib.sha1(s.encode("utf-8")).hexdigest()[:10]


def _edge_witness_fp(e: Any) -> str | None:
    # If stored explicitly, prefer it
    wfp = getattr(e, "witness_fp", None)
    if isinstance(wfp, str) and wfp.strip():
        return wfp
    # Otherwise derive from witness source if present
    w = _edge_witness_src(e)
    return _fp(w) if w else None


def _graph_semantics(g: Any) -> Mapping[str, Any] | None:
    v = getattr(g, "semantics", None)
    return v if isinstance(v, Mapping) else None


def _graph_initials(g: Any) -> list[str]:
    # canonical: initials; tolerate older: initial
    v = getattr(g, "initials", None)
    if v is None:
        v = getattr(g, "initial", None)
    return list(v or [])


def _graph_nodes(g: Any, spec: Any) -> list[str]:
    ns = getattr(g, "nodes", None)
    if ns:
        return [str(x) for x in list(ns) if str(x)]
    # fallback: use spec.scenarios
    out = [getattr(sc, "name", "") for sc in list(getattr(spec, "scenarios", []) or [])]
    return [n for n in out if isinstance(n, str) and n]


def _graph_succ(g: Any, spec: Any) -> dict[str, list[str]]:
    """
    Prefer g.succ if present; otherwise derive from edges.
    """
    succ = getattr(g, "succ", None)
    if isinstance(succ, dict) and succ:
        # normalize to list[str] values
        out: dict[str, list[str]] = {}
        for k, vs in succ.items():
            out[str(k)] = sorted({str(x) for x in (vs or []) if str(x)})
        # Ensure all nodes are present
        for n in _graph_nodes(g, spec):
            out.setdefault(n, [])
        return out

    # fallback: edges -> adjacency
    out: dict[str, list[str]] = {n: [] for n in _graph_nodes(g, spec)}
    for e in _edges(g):
        if not _edge_reachable(e):
            continue
        s, t = _edge_src(e), _edge_dst(e)
        if not s or not t:
            continue
        out.setdefault(s, []).append(t)
    for k in out:
        out[k] = sorted(set(out[k]))
    return out


def _graph_pred(g: Any, spec: Any) -> dict[str, list[str]]:
    """
    Prefer g.pred if present; otherwise derive from succ.
    """
    pred = getattr(g, "pred", None)
    if isinstance(pred, dict) and pred:
        out: dict[str, list[str]] = {}
        for k, vs in pred.items():
            out[str(k)] = sorted({str(x) for x in (vs or []) if str(x)})
        for n in _graph_nodes(g, spec):
            out.setdefault(n, [])
        return out

    succ = _graph_succ(g, spec)
    out: dict[str, list[str]] = {n: [] for n in succ.keys()}
    for s, outs in succ.items():
        for t in outs:
            out.setdefault(t, []).append(s)
    for k in out:
        out[k] = sorted(set(out[k]))
    return out


def _bfs_reachable(initial: Sequence[str], succ: Mapping[str, Sequence[str]]) -> set[str]:
    seen: set[str] = set()
    q: deque[str] = deque()

    for s in initial:
        if s and s not in seen:
            seen.add(s)
            q.append(s)

    while q:
        cur = q.popleft()
        for nxt in succ.get(cur, []):
            if nxt not in seen:
                seen.add(nxt)
                q.append(nxt)

    return seen


def _get_scenario(spec: Any, name: str) -> Any | None:
    for sc in list(getattr(spec, "scenarios", []) or []):
        if getattr(sc, "name", None) == name:
            return sc
    return None


def _scenario_status_reason(sc: Any) -> str | None:
    """
    Best-effort explanation for why a scenario may not participate in reachability.
    """
    if sc is None:
        return "scenario_not_found_in_spec"

    st = getattr(sc, "component_status", None)
    st_name = st.__class__.__name__ if st is not None else ""

    if st_name.lower() == "missing":
        mp = getattr(st, "missing_preds", None) or getattr(sc, "preds_missing", None) or []
        mt = getattr(st, "missing_trans", None) or getattr(sc, "trans_missing", None) or []
        if mp or mt:
            mp = list(mp)
            mt = list(mt)
            return (
                f"missing_components("
                f"preds={mp[:3]}{'...' if len(mp) > 3 else ''}, "
                f"trans={mt[:3]}{'...' if len(mt) > 3 else ''}"
                f")"
            )
        return "missing_components"

    if st_name.lower() == "inconsistent":
        return "inconsistent_predicates"

    # If a scenario has no transitions, it cannot serve as a reachability "step source".
    then = list(getattr(sc, "then", []) or [])
    if len(then) == 0:
        return "no_transitions_in_then"

    return None


def _first_edge_notes(e: Any) -> str | None:
    n = getattr(e, "notes", None)
    return n if isinstance(n, str) and n.strip() else None


# -----------------------------------------------------------------------------
# Public helpers (Spec/state API helpers)
# -----------------------------------------------------------------------------

def list_unreachable_scenarios(spec: Any) -> list[str]:
    g = _graph(spec)
    if g is None:
        return []

    # Prefer precomputed unreachable if present (canonical graph)
    u = getattr(g, "unreachable", None)
    if isinstance(u, list):
        return [str(x) for x in u if str(x)]

    initials = _graph_initials(g)
    if not initials:
        return []

    nodes = _graph_nodes(g, spec)
    succ = _graph_succ(g, spec)
    reached = _bfs_reachable(initials, succ)
    return sorted([n for n in nodes if n not in reached])


def list_successors(spec: Any, scenario_name: str) -> list[str]:
    g = _graph(spec)
    if g is None:
        return []
    succ = _graph_succ(g, spec)
    return list(succ.get(scenario_name, []))


def list_predecessors(spec: Any, scenario_name: str) -> list[str]:
    g = _graph(spec)
    if g is None:
        return []
    pred = _graph_pred(g, spec)
    return list(pred.get(scenario_name, []))


def explain_edge(
    spec: Any,
    src: str,
    dst: str,
    *,
    include_witness_src: bool = True,
    semantics: Mapping[str, Any] | None = None,
) -> EdgeExplanation:
    g = _graph(spec)
    if g is None:
        return EdgeExplanation(reachable=False, notes="scenario_reachability not computed")

    sem = semantics or _graph_semantics(g)

    # Prefer explicit edge record if stored
    for e in _edges(g):
        if _edge_src(e) == src and _edge_dst(e) == dst:
            w = _edge_witness_src(e)
            return EdgeExplanation(
                reachable=_edge_reachable(e),
                witness_fp=_edge_witness_fp(e),
                witness_src=(w if (include_witness_src and w) else None),
                semantics=sem,
                notes=_first_edge_notes(e),
            )

    # Otherwise infer from adjacency
    succ = _graph_succ(g, spec)
    reachable = dst in set(succ.get(src, []))
    return EdgeExplanation(
        reachable=reachable,
        witness_fp=None,
        witness_src=None,
        semantics=sem,
        notes="no explicit edge stored; inferred from adjacency",
    )


def shortest_path(spec: Any, frm: str, to: str) -> list[str] | None:
    """
    BFS shortest path over succ adjacency.
    If the graph model already provides shortest_path(), we delegate to it.
    """
    g = _graph(spec)
    if g is None:
        return None

    # Delegate if the canonical graph now owns this method
    meth = getattr(g, "shortest_path", None)
    if callable(meth):
        try:
            return meth(frm, to)  # type: ignore[misc]
        except TypeError:
            # older signature
            pass

    succ = _graph_succ(g, spec)
    if frm == to:
        return [frm]
    if frm not in succ:
        return None

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

    while q:
        cur = q.popleft()
        for nxt in succ.get(cur, []):
            if nxt in parent:
                continue
            parent[nxt] = cur
            if nxt == to:
                # reconstruct
                path = [to]
                p = cur
                while p is not None:
                    path.append(p)
                    p = parent[p]
                path.reverse()
                return path
            q.append(nxt)

    return None


def frontier_gaps(spec: Any) -> list[tuple[str, str]]:
    """
    Returns list[(src, missing_reason)].

    Heuristic definition of "frontier":
      - if graph has initials => src ranges over reachable-from-initials
      - else => src ranges over all nodes
      - in both cases, src is frontier if it has no outgoing successors
    """
    g = _graph(spec)
    if g is None:
        return []

    # Delegate if the canonical graph now owns this method
    meth = getattr(g, "frontier_gaps", None)
    if callable(meth):
        try:
            return list(meth())  # type: ignore[misc]
        except TypeError:
            pass

    succ = _graph_succ(g, spec)
    initials = _graph_initials(g)

    # Prefer precomputed reachable_from_initials if present
    if initials:
        rfi = getattr(g, "reachable_from_initials", None)
        if isinstance(rfi, list) and rfi:
            frontier = sorted({str(x) for x in rfi if str(x)})
        else:
            frontier = sorted(_bfs_reachable(initials, succ))
    else:
        frontier = sorted(set(_graph_nodes(g, spec)))

    out: list[tuple[str, str]] = []
    for src in frontier:
        if succ.get(src, []):
            continue

        sc = _get_scenario(spec, src)
        reason = _scenario_status_reason(sc)

        if reason is None:
            sem = _graph_semantics(g) or {}
            require_to_when = bool(sem.get("require_to_when", False))
            action_mode = sem.get("action_mode", None)
            extra = []
            if require_to_when:
                extra.append("require_to_when=true")
            if action_mode:
                extra.append(f"action_mode={action_mode}")
            reason = (
                f"no_successors_under_current_semantics({', '.join(extra)})"
                if extra
                else "no_successors_under_current_semantics"
            )

        out.append((src, reason))

    return out