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

from __future__ import annotations

from typing import Any

from .inst import StateInstance
from .diff import (
    StateDiff,
    ValueDiff,
    base_status_comp,
    bool_true_is_good,
    complement_present_comp,
    numeric_increase_bad,
    numeric_increase_good
)
from ..modeling.domain import BaseStatus
from ..modeling.reachability import (
    ScenarioReachabilityGraph,
    ScenarioReachabilityGraphDiff,
    diff_reachability_graph,
)
from ..modeling.trace_reachability import (
    TraceReachabilityAnalysis,
    TraceReachabilityDiff,
    diff_trace_reachability
)


# -----------------------------------------------------------------------------
# Complement helpers (keep tolerant to model evolution)
# -----------------------------------------------------------------------------

def _comp_regions(comp: Any) -> list[Any]:
    """
    Best-effort accessor for complement regions.
    We only need a count here; keep it tolerant to refactors.
    """
    if comp is None:
        return []
    # common shapes you’ve used:
    # - comp.regions_all
    # - comp.regions_preview
    ra = getattr(comp, "regions_all", None)
    if ra:
        return list(ra)
    rp = getattr(comp, "regions_preview", None)
    if rp:
        return list(rp)
    return []


# -----------------------------------------------------------------------------
# Reachability helpers
# -----------------------------------------------------------------------------

def _reach_counts(g: ScenarioReachabilityGraph | None) -> dict[str, int]:
    if g is None:
        return {
            "present": 0,
            "nodes": 0,
            "edges": 0,
            "initials": 0,
            "reachable": 0,
            "unreachable": 0,
        }
    return {
        "present": 1,
        "nodes": len(getattr(g, "nodes", []) or []),
        "edges": len(getattr(g, "edges", []) or []),
        "initials": len(getattr(g, "initials", []) or []),
        "reachable": len(getattr(g, "reachable_from_initials", []) or []),
        "unreachable": len(getattr(g, "unreachable", []) or []),
    }


def _mk_reach_diffs(
    before_g: ScenarioReachabilityGraph | None,
    after_g: ScenarioReachabilityGraph | None,
) -> tuple[
    ValueDiff[bool],
    ValueDiff[int],
    ValueDiff[int],
    ValueDiff[int],
    ValueDiff[int],
    ValueDiff[int],
    ScenarioReachabilityGraphDiff | None,
]:
    b = _reach_counts(before_g)
    a = _reach_counts(after_g)

    # Presence: having a computed graph is generally good (enables queries + agent loop)
    present_vd = ValueDiff(
        label="Reachability graph present",
        before=bool(b["present"]),
        after=bool(a["present"]),
        # If you still prefer "present is good", swap to bool_true_is_good.
        # Presence itself is arguably neutral, but we keep your existing intent.
        comp_func=bool_true_is_good,
        hint={},
    )

    nodes_vd = ValueDiff(
        label="Reachability nodes",
        before=b["nodes"],
        after=a["nodes"],
        comp_func=numeric_increase_good,
        hint={},
    )
    edges_vd = ValueDiff(
        label="Reachability edges",
        before=b["edges"],
        after=a["edges"],
        comp_func=numeric_increase_good,
        hint={},
    )

    reach_initials_vd = ValueDiff(
        label="Initial scenarios (count)",
        before=b["initials"],
        after=a["initials"],
        comp_func=numeric_increase_good,
        hint={},
    )

    reach_reachable_vd = ValueDiff(
        label="Reachable scenarios from initials",
        before=b["reachable"],
        after=a["reachable"],
        comp_func=numeric_increase_good,
        hint={},
    )

    reach_unreachable_vd = ValueDiff(
        label="Unreachable scenarios from initials",
        before=b["unreachable"],
        after=a["unreachable"],
        comp_func=numeric_increase_bad,
        hint={},
    )

    reach_diff = diff_reachability_graph(before_g, after_g) if (before_g or after_g) else None

    return (
        present_vd,
        nodes_vd,
        edges_vd,
        reach_initials_vd,
        reach_reachable_vd,
        reach_unreachable_vd,
        reach_diff,
    )


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

def _tr_counts(a: TraceReachabilityAnalysis | None) -> dict[str, int]:
    if a is None:
        return {
            "present": 0,
            "total_traces": 0,
            "total_steps": 0,
            "total_hops": 0,
            "supported_hops": 0,
            "unsupported_hops": 0,
            "scenarios_hit": 0,
            "scenarios_unhit": 0,
        }

    s = getattr(a, "summary", None)
    if s is None:
        return {
            "present": 1,
            "total_traces": 0,
            "total_steps": 0,
            "total_hops": 0,
            "supported_hops": 0,
            "unsupported_hops": 0,
            "scenarios_hit": 0,
            "scenarios_unhit": 0,
        }

    return {
        "present": 1,
        "total_traces": int(getattr(s, "total_traces", 0) or 0),
        "total_steps": int(getattr(s, "total_steps", 0) or 0),
        "total_hops": int(getattr(s, "total_hops", 0) or 0),
        "supported_hops": int(getattr(s, "supported_hops", 0) or 0),
        "unsupported_hops": int(getattr(s, "unsupported_hops", 0) or 0),
        "scenarios_hit": len(getattr(s, "scenarios_hit", []) or []),
        "scenarios_unhit": len(getattr(s, "scenarios_unhit", []) or []),
    }


def _mk_trace_reach_diffs(
    before_tr: TraceReachabilityAnalysis | None,
    after_tr: TraceReachabilityAnalysis | None,
) -> tuple[
    ValueDiff[bool],
    ValueDiff[int],
    ValueDiff[int],
    ValueDiff[int],
    ValueDiff[int],
    ValueDiff[int],
    ValueDiff[int],
    ValueDiff[int],
    TraceReachabilityDiff | None,
]:
    b = _tr_counts(before_tr)
    a = _tr_counts(after_tr)

    # Presence: neutral (don’t force “good/bad”)
    present_vd = ValueDiff(
        label="Trace reachability present",
        before=bool(b["present"]),
        after=bool(a["present"]),
        comp_func=complement_present_comp,
        hint={},
    )

    total_traces_vd = ValueDiff(
        label="Traces (reachability)",
        before=b["total_traces"],
        after=a["total_traces"],
        comp_func=numeric_increase_good,
        hint={},
    )
    total_steps_vd = ValueDiff(
        label="Steps (reachability)",
        before=b["total_steps"],
        after=a["total_steps"],
        comp_func=numeric_increase_good,
        hint={},
    )
    total_hops_vd = ValueDiff(
        label="Hops (reachability)",
        before=b["total_hops"],
        after=a["total_hops"],
        comp_func=numeric_increase_good,
        hint={},
    )

    supported_hops_vd = ValueDiff(
        label="Supported hops",
        before=b["supported_hops"],
        after=a["supported_hops"],
        comp_func=numeric_increase_good,
        hint={},
    )
    unsupported_hops_vd = ValueDiff(
        label="Unsupported hops",
        before=b["unsupported_hops"],
        after=a["unsupported_hops"],
        comp_func=numeric_increase_bad,
        hint={},
    )

    scenarios_hit_vd = ValueDiff(
        label="Scenarios hit (count)",
        before=b["scenarios_hit"],
        after=a["scenarios_hit"],
        comp_func=numeric_increase_good,
        hint={},
    )
    scenarios_unhit_vd = ValueDiff(
        label="Scenarios unhit (count)",
        before=b["scenarios_unhit"],
        after=a["scenarios_unhit"],
        comp_func=numeric_increase_bad,
        hint={},
    )

    struct = diff_trace_reachability(before_tr, after_tr) if (before_tr or after_tr) else None

    return (
        present_vd,
        total_traces_vd,
        total_steps_vd,
        total_hops_vd,
        supported_hops_vd,
        unsupported_hops_vd,
        scenarios_hit_vd,
        scenarios_unhit_vd,
        struct,
    )


# -----------------------------------------------------------------------------
# Main diff
# -----------------------------------------------------------------------------

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

    Assumes `old` and `new` have fully computed stats on:
      - spec.domain_model
      - spec (scenario stats, conflicts, complement, reachability if enabled)
      - art_container
      - trace_reachability (if enabled)
    """

    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

    state_diff_fields = getattr(StateDiff, "model_fields", {})

    out_kwargs: dict[str, Any] = {}

    def _emit(key: str, vd: ValueDiff) -> None:
        if key in state_diff_fields:
            out_kwargs[key] = vd

    def _set(key: str, value: Any) -> None:
        if key in state_diff_fields:
            out_kwargs[key] = value

    # -------------------------------------------------------------------------
    # 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={},
        ),
    )

    # -------------------------------------------------------------------------
    # Predicates / transitions (domain model counters)
    # -------------------------------------------------------------------------
    _emit("num_state_preds_total", ValueDiff(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(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(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(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={}))

    _emit("num_action_preds_total", ValueDiff(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(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(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(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={}))

    _emit("num_preds_total", ValueDiff(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(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(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(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={}))

    _emit("num_trans_total", ValueDiff(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(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(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(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 / conflicts
    # -------------------------------------------------------------------------
    _emit("num_sc_total", ValueDiff(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(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(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(label="Number of inconsistent scenarios", before=spec_old.num_sc_inconsistent, after=spec_new.num_sc_inconsistent, comp_func=numeric_increase_bad, hint={}))

    _emit("num_sc_conflicted", ValueDiff(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(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(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
    # -------------------------------------------------------------------------
    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)

    _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={},
        ),
    )
    _emit(
        "num_comp_regions_total",
        ValueDiff[int | None](
            label="Complement region count",
            before=(len(old_regions) if comp_old is not None else None),
            after=(len(new_regions) if comp_new is not None else None),
            # Smaller complement = better coverage; reuse numeric_increase_bad by mapping “increase = bad”.
            # When None appears, treat as UNKNOWN via ValueDiff.is_different fallback; renderer can handle.
            comp_func=None,
            hint={},
        ),
    )

    # -------------------------------------------------------------------------
    # Traces
    # -------------------------------------------------------------------------
    _emit("num_test_traces_total", ValueDiff(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(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(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(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(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(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={}))

    # -------------------------------------------------------------------------
    # Reachability graph (spec-level)
    # -------------------------------------------------------------------------
    before_g = getattr(spec_old, "scenario_reachability", None)
    after_g = getattr(spec_new, "scenario_reachability", None)

    (
        reach_present_vd,
        reach_nodes_vd,
        reach_edges_vd,
        reach_initials_vd,
        reach_reachable_vd,
        reach_unreachable_vd,
        reach_struct_diff,
    ) = _mk_reach_diffs(before_g, after_g)

    _emit("scenario_reach_present", reach_present_vd)
    _emit("reach_num_nodes", reach_nodes_vd)
    _emit("reach_num_edges", reach_edges_vd)
    _emit("reach_num_initials", reach_initials_vd)
    _emit("reach_num_reachable", reach_reachable_vd)
    _emit("reach_num_unreachable", reach_unreachable_vd)
    _set("scenario_reach_diff", reach_struct_diff)

    # -------------------------------------------------------------------------
    # Trace reachability analysis (state-level)
    # -------------------------------------------------------------------------
    before_tr = getattr(old, "trace_reachability", None)
    after_tr = getattr(new, "trace_reachability", None)

    (
        tr_present_vd,
        tr_total_traces_vd,
        tr_total_steps_vd,
        tr_total_hops_vd,
        tr_supported_hops_vd,
        tr_unsupported_hops_vd,
        tr_scenarios_hit_vd,
        tr_scenarios_unhit_vd,
        tr_struct_diff,
    ) = _mk_trace_reach_diffs(before_tr, after_tr)

    _emit("trace_reach_present", tr_present_vd)
    _emit("tr_total_traces", tr_total_traces_vd)
    _emit("tr_total_steps", tr_total_steps_vd)
    _emit("tr_total_hops", tr_total_hops_vd)
    _emit("tr_supported_hops", tr_supported_hops_vd)
    _emit("tr_unsupported_hops", tr_unsupported_hops_vd)
    _emit("tr_scenarios_hit", tr_scenarios_hit_vd)
    _emit("tr_scenarios_unhit", tr_scenarios_unhit_vd)
    _set("trace_reach_diff", tr_struct_diff)

    return StateDiff(**out_kwargs)