#
#   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 .change import StateChange
from .diff import (
    StateDiff,
    ValueDiff,
    base_status_comp,
    bool_true_is_good,
    numeric_increase_bad,
    numeric_increase_good,
)
from .recommendation import Recommendation
from ..data.container import ArtifactContainer
from ..data.mapping import ArtifactMap
from ..modeling.domain import BaseStatus
from ..modeling.spec import Spec
from ..utils import JSONObject

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])

    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]


def calc_diff(new: StateInstance, old: StateInstance) -> StateDiff:
    """
    Generate a diff between two StateInstances.

    Assumes `old` and `new` have fully computed stats on:
      - new.spec.domain_model
      - new.spec (scenario stats)
      - new.art_container
      - new.spec.scenario_comp (optional)
    """
    vals: list[tuple[str, ValueDiff]] = []

    dm_old = old.spec.domain_model
    dm_new = new.spec.domain_model
    spec_old = old.spec
    spec_new = new.spec
    art_old = old.art_container
    art_new = new.art_container

    # -------------------------------------------------------------------------
    # Helper: only emit diffs for fields that actually exist on StateDiff.
    # This makes calc_diff resilient while you evolve StateDiff.
    # -------------------------------------------------------------------------
    state_diff_fields = getattr(StateDiff, "model_fields", {})  # pydantic v2

    def _emit(key: str, vd: ValueDiff) -> None:
        if key in state_diff_fields:
            vals.append((key, vd))

    # >>> Base
    _emit(
        "base_status",
        ValueDiff[BaseStatus](
            label="Base code IML status",
            before=dm_old.base_status,
            after=dm_new.base_status,
            comp_func=base_status_comp,
            hint={},
        ),
    )

    _emit(
        "base_has_state",
        ValueDiff[bool](
            label="State type is present",
            before=dm_old.base_has_state,
            after=dm_new.base_has_state,
            comp_func=bool_true_is_good,
            hint={},
        ),
    )

    _emit(
        "base_has_action",
        ValueDiff[bool](
            label="Action type is present",
            before=dm_old.base_has_action,
            after=dm_new.base_has_action,
            comp_func=bool_true_is_good,
            hint={},
        ),
    )

    # >>> State predicates
    _emit(
        "num_state_preds_total",
        ValueDiff[int](
            label="Number of state predicates",
            before=dm_old.num_state_preds_total,
            after=dm_new.num_state_preds_total,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_state_preds_matched",
        ValueDiff[int](
            label="Number of matched state predicates",
            before=dm_old.num_state_preds_matched,
            after=dm_new.num_state_preds_matched,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_state_preds_valid_logic",
        ValueDiff[int](
            label="Number of state predicates with valid logic",
            before=dm_old.num_state_preds_valid_logic,
            after=dm_new.num_state_preds_valid_logic,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_state_preds_errored",
        ValueDiff[int](
            label="Number of state predicates with errors",
            before=dm_old.num_state_preds_errored,
            after=dm_new.num_state_preds_errored,
            comp_func=numeric_increase_bad,
            hint={},
        ),
    )

    # >>> Action predicates
    _emit(
        "num_action_preds_total",
        ValueDiff[int](
            label="Number of action predicates",
            before=dm_old.num_action_preds_total,
            after=dm_new.num_action_preds_total,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_action_preds_matched",
        ValueDiff[int](
            label="Number of matched action predicates",
            before=dm_old.num_action_preds_matched,
            after=dm_new.num_action_preds_matched,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_action_preds_valid_logic",
        ValueDiff[int](
            label="Number of action predicates with valid logic",
            before=dm_old.num_action_preds_valid_logic,
            after=dm_new.num_action_preds_valid_logic,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_action_preds_errored",
        ValueDiff[int](
            label="Number of action predicates with errors",
            before=dm_old.num_action_preds_errored,
            after=dm_new.num_action_preds_errored,
            comp_func=numeric_increase_bad,
            hint={},
        ),
    )

    # >>> Predicates (totals)
    _emit(
        "num_preds_total",
        ValueDiff[int](
            label="Number of predicates (total)",
            before=dm_old.num_preds_total,
            after=dm_new.num_preds_total,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_preds_matched",
        ValueDiff[int](
            label="Number of matched predicates (total)",
            before=dm_old.num_preds_matched,
            after=dm_new.num_preds_matched,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_preds_valid_logic",
        ValueDiff[int](
            label="Number of predicates with valid logic (total)",
            before=dm_old.num_preds_valid_logic,
            after=dm_new.num_preds_valid_logic,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_preds_errored",
        ValueDiff[int](
            label="Number of predicates with errors (total)",
            before=dm_old.num_preds_errored,
            after=dm_new.num_preds_errored,
            comp_func=numeric_increase_bad,
            hint={},
        ),
    )

    # >>> Transitions
    _emit(
        "num_trans_total",
        ValueDiff[int](
            label="Number of transitions",
            before=dm_old.num_trans_total,
            after=dm_new.num_trans_total,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_trans_matched",
        ValueDiff[int](
            label="Number of matched transitions",
            before=dm_old.num_trans_matched,
            after=dm_new.num_trans_matched,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_trans_valid_logic",
        ValueDiff[int](
            label="Number of transitions with valid logic",
            before=dm_old.num_trans_valid_logic,
            after=dm_new.num_trans_valid_logic,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_trans_errored",
        ValueDiff[int](
            label="Number of transitions with errors",
            before=dm_old.num_trans_errored,
            after=dm_new.num_trans_errored,
            comp_func=numeric_increase_bad,
            hint={},
        ),
    )

    # >>> Scenarios
    _emit(
        "num_sc_total",
        ValueDiff[int](
            label="Number of scenarios",
            before=spec_old.num_sc_total,
            after=spec_new.num_sc_total,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_sc_missing",
        ValueDiff[int](
            label="Number of scenarios with missing model components",
            before=spec_old.num_sc_missing,
            after=spec_new.num_sc_missing,
            comp_func=numeric_increase_bad,
            hint={},
        ),
    )

    _emit(
        "num_sc_matched",
        ValueDiff[int](
            label="Number of scenarios matched",
            before=spec_old.num_sc_matched,
            after=spec_new.num_sc_matched,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_sc_inconsistent",
        ValueDiff[int](
            label="Number of inconsistent scenarios",
            before=spec_old.num_sc_inconsistent,
            after=spec_new.num_sc_inconsistent,
            comp_func=numeric_increase_bad,
            hint={},
        ),
    )

    # >>> Conflicts
    _emit(
        "num_sc_conflicted",
        ValueDiff[int](
            label="Number of scenario conflicts (total)",
            before=spec_old.num_sc_conflicted,
            after=spec_new.num_sc_conflicted,
            comp_func=numeric_increase_bad,
            hint={},
        ),
    )

    _emit(
        "num_sc_overlap",
        ValueDiff[int](
            label="Number of scenario overlap conflicts",
            before=getattr(spec_old, "num_sc_overlap", 0),
            after=getattr(spec_new, "num_sc_overlap", 0),
            comp_func=numeric_increase_bad,
            hint={},
        ),
    )

    _emit(
        "num_sc_consumed",
        ValueDiff[int](
            label="Number of scenario consumed conflicts",
            before=getattr(spec_old, "num_sc_consumed", 0),
            after=getattr(spec_new, "num_sc_consumed", 0),
            comp_func=numeric_increase_bad,
            hint={},
        ),
    )

    # -------------------------------------------------------------------------
    # >>> Complement (Spec-level): scenario_comp
    #   NOTE: match the new StateDiff field names you mentioned:
    #     - scenario_comp_present
    #     - num_comp_regions_total
    # -------------------------------------------------------------------------
    comp_old = getattr(spec_old, "scenario_comp", None)
    comp_new = getattr(spec_new, "scenario_comp", None)

    old_regions = _comp_regions(comp_old)
    new_regions = _comp_regions(comp_new)

    # Presence: having a computed complement is generally "good" (enables coverage reasoning).
    _emit(
        "scenario_comp_present",
        ValueDiff[bool](
            label="Scenario complement is present",
            before=(comp_old is not None),
            after=(comp_new is not None),
            comp_func=bool_true_is_good,
            hint={},
        ),
    )

    # Region count: smaller complement is better (more coverage), so increases are "bad".
    _emit(
        "num_comp_regions_total",
        ValueDiff[int](
            label="Complement region count",
            before=len(old_regions),
            after=len(new_regions),
            comp_func=numeric_increase_bad,
            hint={},
        ),
    )

    # -------------------------------------------------------------------------
    # >>> Traces
    # -------------------------------------------------------------------------
    _emit(
        "num_test_traces_total",
        ValueDiff[int](
            label="Number of test traces",
            before=art_old.num_test_traces_total,
            after=art_new.num_test_traces_total,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_test_traces_logic_good",
        ValueDiff[int](
            label="Number of test traces logically valid",
            before=art_old.num_test_traces_logic_good,
            after=art_new.num_test_traces_logic_good,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_test_traces_matched",
        ValueDiff[int](
            label="Number of test traces matched",
            before=art_old.num_test_traces_matched,
            after=art_new.num_test_traces_matched,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_log_traces_total",
        ValueDiff[int](
            label="Number of log traces",
            before=art_old.num_log_traces_total,
            after=art_new.num_log_traces_total,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_log_traces_logic_good",
        ValueDiff[int](
            label="Number of log traces logically valid",
            before=art_old.num_log_traces_logic_good,
            after=art_new.num_log_traces_logic_good,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    _emit(
        "num_log_traces_matched",
        ValueDiff[int](
            label="Number of log traces matched",
            before=art_old.num_log_traces_matched,
            after=art_new.num_log_traces_matched,
            comp_func=numeric_increase_good,
            hint={},
        ),
    )

    return StateDiff(**dict(vals))