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

from pydantic import BaseModel, Field
from typing import Literal

class ObservedScenarioHop(BaseModel):
    trace_id: str
    step_from: int
    step_to: int
    src: str
    dst: str
    supported_by_graph: bool
    shortest_path: list[str] | None = None   # if unsupported, try to explain
    reason: str | None = None                # "no_graph_edge", "no_path", etc.

class TraceScenarioAlignment(BaseModel):
    trace_id: str
    # scenario candidates per step (could be large; keep optional or capped)
    step_matches: list[list[str]] = Field(default_factory=list)
    # chosen label per step (optional, if you implement ranking)
    step_best: list[str | None] = Field(default_factory=list)

    # derived hops between consecutive steps
    observed_hops: list[ObservedScenarioHop] = Field(default_factory=list)

    # quick stats for UI
    num_steps: int = 0
    num_matched_steps: int = 0
    num_supported_hops: int = 0
    num_unsupported_hops: int = 0

class TraceReachabilitySummary(BaseModel):
    total_traces: int = 0
    total_steps: int = 0
    total_hops: int = 0
    supported_hops: int = 0
    unsupported_hops: int = 0

    # scenario coverage
    scenarios_hit: list[str] = Field(default_factory=list)
    scenarios_unhit: list[str] = Field(default_factory=list)

    # ranked issues
    top_unsupported_edges: list[tuple[str, str, int]] = Field(default_factory=list)  # (src,dst,count)
    top_frontier_trace_endings: list[tuple[str, int]] = Field(default_factory=list) # (scenario,count)

class TraceReachabilityAnalysis(BaseModel):
    kind: Literal["trace_reachability"] = "trace_reachability"

    # snapshot fingerprints so you can detect stale analyses
    spec_fp: str | None = None
    artifacts_fp: str | None = None
    mapping_fp: str | None = None

    per_trace: dict[str, TraceScenarioAlignment] = Field(default_factory=dict)
    summary: TraceReachabilitySummary = Field(default_factory=TraceReachabilitySummary)


class TraceReachabilityDiff(BaseModel):
    """
    Deterministic structural diff for trace reachability analysis.

    Notes:
      - We treat top_* lists as sets of keys to avoid ordering noise.
      - We diff scenario coverage as sets.
    """

    present_before: bool = False
    present_after: bool = False

    scenarios_hit_added: list[str] = Field(default_factory=list)
    scenarios_hit_removed: list[str] = Field(default_factory=list)

    scenarios_unhit_added: list[str] = Field(default_factory=list)
    scenarios_unhit_removed: list[str] = Field(default_factory=list)

    # Ranked issue keys (ignore counts for “structural” changes, but keep list deterministic)
    unsupported_edge_keys_added: list[tuple[str, str]] = Field(default_factory=list)
    unsupported_edge_keys_removed: list[tuple[str, str]] = Field(default_factory=list)

    frontier_ending_keys_added: list[str] = Field(default_factory=list)
    frontier_ending_keys_removed: list[str] = Field(default_factory=list)

    def has_changes(self) -> bool:
        return any(
            [
                self.present_before != self.present_after,
                self.scenarios_hit_added,
                self.scenarios_hit_removed,
                self.scenarios_unhit_added,
                self.scenarios_unhit_removed,
                self.unsupported_edge_keys_added,
                self.unsupported_edge_keys_removed,
                self.frontier_ending_keys_added,
                self.frontier_ending_keys_removed,
            ]
        )


def diff_trace_reachability(
    before: TraceReachabilityAnalysis | None,
    after: TraceReachabilityAnalysis | None,
) -> TraceReachabilityDiff:
    present_before = before is not None
    present_after = after is not None

    b_hit = set((before.summary.scenarios_hit if before else []) or [])
    a_hit = set((after.summary.scenarios_hit if after else []) or [])

    b_unhit = set((before.summary.scenarios_unhit if before else []) or [])
    a_unhit = set((after.summary.scenarios_unhit if after else []) or [])

    # keys only (ignore counts so it’s stable)
    def _edge_keys(a: TraceReachabilityAnalysis | None) -> set[tuple[str, str]]:
        xs = (a.summary.top_unsupported_edges if a else []) or []
        return {(src, dst) for (src, dst, _n) in xs}

    def _frontier_keys(a: TraceReachabilityAnalysis | None) -> set[str]:
        xs = (a.summary.top_frontier_trace_endings if a else []) or []
        return {sc for (sc, _n) in xs}

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

    b_frontier = _frontier_keys(before)
    a_frontier = _frontier_keys(after)

    return TraceReachabilityDiff(
        present_before=present_before,
        present_after=present_after,
        scenarios_hit_added=sorted(a_hit - b_hit),
        scenarios_hit_removed=sorted(b_hit - a_hit),
        scenarios_unhit_added=sorted(a_unhit - b_unhit),
        scenarios_unhit_removed=sorted(b_unhit - a_unhit),
        unsupported_edge_keys_added=sorted(a_edges - b_edges),
        unsupported_edge_keys_removed=sorted(b_edges - a_edges),
        frontier_ending_keys_added=sorted(a_frontier - b_frontier),
        frontier_ending_keys_removed=sorted(b_frontier - a_frontier),
    )
