#
#   Imandra Inc.
#
#   speclogician/state/inst.py
#

from __future__ import annotations

from datetime import datetime, timezone
from typing import Optional

from pydantic import BaseModel, Field

from ..data.container import ArtifactContainer
from ..data.mapping import ArtifactMap
from ..modeling.spec import Spec
from ..modeling.trace_reachability import (
    TraceReachabilityAnalysis, 
    TraceScenarioAlignment, 
    ObservedScenarioHop
)
from ..utils import JSONObject

from .change import StateChange
from .diff import StateDiff
from .recommendation import Recommendation

def _safe_iso(dt: Optional[datetime]) -> str | None:
    if dt is None:
        return None
    return dt.isoformat()

def _comp_regions(comp: object | None) -> list[object]:
    """
    Best-effort extraction of complement regions.

    Supports a couple of plausible layouts:
      - comp.regions
      - comp.cover_regions
    """
    if comp is None:
        return []
    regs = getattr(comp, "regions", None)
    if regs is None:
        regs = getattr(comp, "cover_regions", None)
    try:
        return list(regs or [])
    except Exception:
        return []


class StateInstance(BaseModel):
    """
    State Instance class
    """

    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))

    spec: Spec = Field(default_factory=Spec)  # `Spec` is the domain model + scenarios
    art_container: ArtifactContainer = Field(default_factory=ArtifactContainer)
    art_map: ArtifactMap = Field(default_factory=ArtifactMap)

    # These describe the changes from the previous state
    changes: list[StateChange] = Field(default_factory=list[StateChange])
    state_diff: StateDiff | None = None

    # The generated recommendations
    recommendations: list[Recommendation] = Field(default_factory=list[Recommendation])

    # Trace reachability analysis
    trace_reachability: TraceReachabilityAnalysis | None = None

    # ----------------------------
    # Trace reachability helpers
    # ----------------------------

    def get_trace_alignment(self, trace_id: str) -> TraceScenarioAlignment | None:
        ana = self.trace_reachability
        if ana is None:
            return None
        return (ana.per_trace or {}).get(trace_id)

    def get_trace_hops(self, trace_id: str) -> list[ObservedScenarioHop]:
        a = self.get_trace_alignment(trace_id)
        return list(a.observed_hops) if a is not None and a.observed_hops else []

    def get_trace_best_path(self, trace_id: str) -> list[str | None]:
        a = self.get_trace_alignment(trace_id)
        return list(a.step_best) if a is not None and a.step_best else []

    def get_trace_unsupported_hops(self, trace_id: str) -> list[ObservedScenarioHop]:
        hops = self.get_trace_hops(trace_id)
        return [h for h in hops if not h.supported_by_graph]

    # Optional convenience: last matched scenario (useful for “frontier ending” UI)
    def get_trace_last_matched(self, trace_id: str) -> str | None:
        best = self.get_trace_best_path(trace_id)
        for x in reversed(best):
            if x is not None:
                return x
        return None

    # Optional convenience: list of trace ids in analysis (stable ordering)
    def get_trace_ids_with_alignment(self) -> list[str]:
        ana = self.trace_reachability
        if ana is None or not ana.per_trace:
            return []
        return sorted(ana.per_trace.keys())

    def update_created_at(self) -> None:
        """Update the `created_at` timestamp field"""
        self.created_at = datetime.now(timezone.utc)

    def json_summary(self) -> JSONObject:
        """
        JSON snapshot of the *inputs* used by calc_diff() to compute StateDiff.

        Goal: stable, machine-readable summary suitable for --json output and tests.
        """
        dm = self.spec.domain_model
        spec = self.spec
        art = self.art_container

        comp = getattr(spec, "scenario_comp", None)
        regions = _comp_regions(comp)

        return {
            "created_at": _safe_iso(self.created_at),
            "changes": [c.model_dump(mode="json") for c in (self.changes or [])],
            # ---- Domain model (base) ----
            "base_status": getattr(dm, "base_status", None),
            "base_has_state": getattr(dm, "base_has_state", None),
            "base_has_action": getattr(dm, "base_has_action", None),
            # ---- Domain model (state predicates) ----
            "num_state_preds_total": getattr(dm, "num_state_preds_total", None),
            "num_state_preds_matched": getattr(dm, "num_state_preds_matched", None),
            "num_state_preds_valid_logic": getattr(dm, "num_state_preds_valid_logic", None),
            "num_state_preds_errored": getattr(dm, "num_state_preds_errored", None),
            # ---- Domain model (action predicates) ----
            "num_action_preds_total": getattr(dm, "num_action_preds_total", None),
            "num_action_preds_matched": getattr(dm, "num_action_preds_matched", None),
            "num_action_preds_valid_logic": getattr(dm, "num_action_preds_valid_logic", None),
            "num_action_preds_errored": getattr(dm, "num_action_preds_errored", None),
            # ---- Domain model (predicates totals) ----
            "num_preds_total": getattr(dm, "num_preds_total", None),
            "num_preds_matched": getattr(dm, "num_preds_matched", None),
            "num_preds_valid_logic": getattr(dm, "num_preds_valid_logic", None),
            "num_preds_errored": getattr(dm, "num_preds_errored", None),
            # ---- Domain model (transitions) ----
            "num_trans_total": getattr(dm, "num_trans_total", None),
            "num_trans_matched": getattr(dm, "num_trans_matched", None),
            "num_trans_valid_logic": getattr(dm, "num_trans_valid_logic", None),
            "num_trans_errored": getattr(dm, "num_trans_errored", None),
            # ---- Scenarios (Spec-level) ----
            "num_sc_total": getattr(spec, "num_sc_total", None),
            "num_sc_missing": getattr(spec, "num_sc_missing", None),
            "num_sc_matched": getattr(spec, "num_sc_matched", None),
            "num_sc_inconsistent": getattr(spec, "num_sc_inconsistent", None),
            # Conflicts (total + by kind).
            "num_sc_conflicted": getattr(spec, "num_sc_conflicted", None),
            "num_sc_overlap": getattr(spec, "num_sc_overlap", 0),
            "num_sc_consumed": getattr(spec, "num_sc_consumed", 0),
            # ---- Complement (Spec-level) ----
            # NOTE: match StateDiff field names (new ones).
            "scenario_comp_present": comp is not None,
            "num_comp_regions_total": len(regions),
            # ---- ArtifactContainer: test traces ----
            "num_test_traces_total": getattr(art, "num_test_traces_total", None),
            "num_test_traces_logic_good": getattr(art, "num_test_traces_logic_good", None),
            "num_test_traces_matched": getattr(art, "num_test_traces_matched", None),
            # ---- ArtifactContainer: log traces ----
            "num_log_traces_total": getattr(art, "num_log_traces_total", None),
            "num_log_traces_logic_good": getattr(art, "num_log_traces_logic_good", None),
            "num_log_traces_matched": getattr(art, "num_log_traces_matched", None),
        }

class InstancesList(BaseModel):
    """A helper container for pretty printing a list of instances"""

    states: list[StateInstance]