#
#   Imandra Inc.
#
#   speclogician/logic/lib/scenarios.py
#

from __future__ import annotations

import asyncio
from typing import Iterable

from speclogician.logic.api import SpecLogicianImandraX
from speclogician.modeling.scenario import Scenario, Valid, Inconsistent, Missing
from speclogician.modeling.conflict import (
    ScenarioConflict,
    ScenarioOverlap,
    ScenarioConsumed,
)
from speclogician.modeling.spec import Spec
from speclogician.utils import console


def _select_scenarios(
    spec: Spec, 
    *, 
    names: list[str], 
    predicates: list[str], 
    transitions: list[str]
) -> list[Scenario]:
    name_set = set(names) if names else None
    pred_set = set(predicates) if predicates else None
    trans_set = set(transitions) if transitions else None

    def _name_ok(sc: Scenario) -> bool:
        return (name_set is None) or (sc.name in name_set)

    def _contains_any_pred(sc: Scenario) -> bool:
        if pred_set is None:
            return True
        sc_preds = set(list(sc.given) + list(sc.when))
        return any(p in sc_preds for p in pred_set)

    def _contains_any_trans(sc: Scenario) -> bool:
        if trans_set is None:
            return True
        sc_trans = set(list(sc.then))
        return any(t in sc_trans for t in trans_set)

    return [
        sc
        for sc in spec.scenarios
        if _name_ok(sc) and _contains_any_pred(sc) and _contains_any_trans(sc)
    ]


def _is_checkable(sc: Scenario) -> bool:
    # must not be missing AND must be solver-consistent
    return (not isinstance(sc.component_status, Missing)) and bool(sc.all_preds_consistent)


async def _analyze_one_scenario(
    ix: SpecLogicianImandraX, 
    *, 
    spec: Spec, 
    sc: Scenario
) -> None:
    dm = spec.domain_model

    # reset per-run fields
    sc.preds_missing = []
    sc.trans_missing = []
    sc.given_preds_consistent = True
    sc.when_preds_consistent = True
    sc.all_preds_consistent = True
    sc.component_status = Valid()

    # Missing components
    missing_preds: list[str] = []
    for p in list(sc.given) + list(sc.when):
        if dm.get_pred_by_name(p) is None:
            missing_preds.append(p)

    missing_trans: list[str] = []
    for t in list(sc.then):
        if dm.get_trans_by_name(t) is None:
            missing_trans.append(t)

    sc.preds_missing = missing_preds
    sc.trans_missing = missing_trans

    if missing_preds or missing_trans:
        sc.component_status = Missing(
            missing_preds=missing_preds,
            missing_trans=missing_trans,
        )
        sc.given_preds_consistent = False if missing_preds else True
        sc.when_preds_consistent = False if missing_preds else True
        sc.all_preds_consistent = False
        return

    # local validation (raises on wrong kind / missing)
    dm.assert_has_predicates(given=sc.given, when=sc.when)
    dm.assert_has_transitions(then=sc.then)

    try:
        if list(sc.given):
            r_given = await ix.check_scenario_consistent(
                domain_model=dm,
                scenario=sc,
                check_state_preds=True,
                check_action_preds=False,
            )
            sc.given_preds_consistent = bool(r_given.is_consistent)
        else:
            sc.given_preds_consistent = True

        if list(sc.when):
            r_when = await ix.check_scenario_consistent(
                domain_model=dm,
                scenario=sc,
                check_state_preds=False,
                check_action_preds=True,
            )
            sc.when_preds_consistent = bool(r_when.is_consistent)
        else:
            sc.when_preds_consistent = True

        if list(sc.given) or list(sc.when):
            r_all = await ix.check_scenario_consistent(
                domain_model=dm,
                scenario=sc,
                check_state_preds=True,
                check_action_preds=True,
            )
            sc.all_preds_consistent = bool(r_all.is_consistent)
        else:
            # No predicates at all => trivially consistent
            sc.all_preds_consistent = True

    except Exception:
        # backend/compile errors => treat as inconsistent
        sc.given_preds_consistent = False
        sc.when_preds_consistent = False
        sc.all_preds_consistent = False

    sc.component_status = Valid() if sc.all_preds_consistent else Inconsistent()


async def _pairwise_overlaps(
    ix: SpecLogicianImandraX,
    *, 
    spec: Spec, 
    scenarios: list[Scenario]
) -> list[ScenarioOverlap]:
    dm = spec.domain_model
    checkable = [sc for sc in scenarios if _is_checkable(sc)]
    out: list[ScenarioOverlap] = []

    for i in range(len(checkable)):
        for j in range(i + 1, len(checkable)):
            a = checkable[i]
            b = checkable[j]

            inter = await ix.check_scenarios_intersect(domain_model=dm, sc_a=a, sc_b=b)
            if inter.do_they_intersect:
                witness = getattr(inter, "model_src", None)
                details = "Scenarios overlap (intersection is satisfiable)."
                if witness:
                    details += f"\n\nWitness model:\n{witness}"

                out.append(
                    ScenarioOverlap(
                        scenario_name1=a.name,
                        scenario_name2=b.name,
                        details=details,
                    )
                )

    return out


async def _pairwise_consumed(
    ix: SpecLogicianImandraX, 
    *, 
    spec: Spec, 
    scenarios: list[Scenario]
) -> list[ScenarioConsumed]:
    """
    Compute directional "consumed" conflicts.

    Requires ix.check_scenario_implies(domain_model=..., sc_a=..., sc_b=...)
      which should return:
        - ScenarioConsumed(...) when B => A
        - None otherwise

    Convention used here:
      ScenarioConsumed(scenario_name1=A, scenario_name2=B) means B => A (B is subset of A).
    """
    dm = spec.domain_model
    if not hasattr(ix, "check_scenario_implies"):
        return []

    checkable = [sc for sc in scenarios if _is_checkable(sc)]
    out: list[ScenarioConsumed] = []

    for i in range(len(checkable)):
        for j in range(len(checkable)):
            if i == j:
                continue

            a = checkable[i]  # A (more general)
            b = checkable[j]  # B (more specific)

            # Conservative: only compare "consumption" when transitions are identical.
            # (Predicate implication alone doesn't guarantee identical post-state behavior.)
            if list(a.then) != list(b.then):
                continue

            # Expect: ScenarioConsumed | None
            c = await ix.check_scenario_implies(domain_model=dm, sc_a=a, sc_b=b)
            if c is not None:
                out.append(c)

    return out


def _dedup_conflicts(conflicts: Iterable[ScenarioConflict]) -> list[ScenarioConflict]:
    seen: set[tuple[str, str, str]] = set()
    out: list[ScenarioConflict] = []

    for c in conflicts:
        if isinstance(c, ScenarioOverlap):
            a, b = sorted([c.scenario_name1, c.scenario_name2])
            key = ("overlap", a, b)
        else:
            # consumed is directional (A,B)
            key = ("consumed", c.scenario_name1, c.scenario_name2)

        if key in seen:
            continue
        seen.add(key)
        out.append(c)

    return out


def _refresh_scenario_stats(spec: Spec) -> None:
    """
    Recompute Spec-level scenario stats from the current scenario objects.

    NOTE: This is GLOBAL across spec.scenarios (not just the selected subset).
    """
    from speclogician.modeling.scenario import Missing, Inconsistent

    spec.num_sc_total = len(spec.scenarios)

    spec.num_sc_missing = sum(
        1 for sc in spec.scenarios
        if isinstance(getattr(sc, "component_status", None), Missing)
    )

    spec.num_sc_inconsistent = sum(
        1 for sc in spec.scenarios
        if isinstance(getattr(sc, "component_status", None), Inconsistent)
    )

    # "matched" is typically filled by your trace-matching pass; keep it robust.
    # Prefer a field if you have it (e.g. sc.is_matched / sc.matched), otherwise 0.
    def _is_matched(sc) -> bool:
        v = getattr(sc, "is_matched", None)
        if v is not None:
            return bool(v)
        v = getattr(sc, "matched", None)
        if v is not None:
            return bool(v)
        return False

    spec.num_sc_matched = sum(1 for sc in spec.scenarios if _is_matched(sc))

    # ---- Conflicts (spec-level) ----
    conflicts = list(getattr(spec, "scenario_conflicts", []) or [])
    spec.num_sc_conflicted = len(conflicts)

    # Count by kind (be tolerant if classes change)
    from speclogician.modeling.conflict import ScenarioOverlap, ScenarioConsumed

    spec.num_sc_overlap = sum(1 for c in conflicts if isinstance(c, ScenarioOverlap))
    spec.num_sc_consumed = sum(1 for c in conflicts if isinstance(c, ScenarioConsumed))

def check_scenarios(
    spec: Spec,
    names: list[str] = [],
    predicates: list[str] = [],
    transitions: list[str] = [],
    json_only: bool = False,
    *,
    analyze_conflicts: bool = True,
    analyze_consumed: bool = True,
) -> None:
    """
    Analyze scenarios:
      - Filter scenarios by name/predicate/transition (conjunctive filters).
      - For each selected scenario:
          * set preds_missing/trans_missing
          * set component_status (Missing/Valid/Inconsistent)
          * set given_preds_consistent/when_preds_consistent/all_preds_consistent
      - Optionally compute conflicts at Spec level:
          * overlaps via intersection
          * consumed via implication (if ix.check_scenario_implies exists)

    Note: conflicts are stored on spec.scenario_conflicts.
    """
    selected = _select_scenarios(
        spec, names=names, predicates=predicates, transitions=transitions
    )

    if not json_only:
        console.print(f"[bold]Logician[/bold] Checking scenarios: {len(selected)} selected")

    async def _run() -> None:
        async with SpecLogicianImandraX() as ix:
            # per-scenario checks
            for sc in selected:
                await _analyze_one_scenario(ix, spec=spec, sc=sc)

            if not analyze_conflicts:
                return

            overlaps = await _pairwise_overlaps(ix, spec=spec, scenarios=selected)
            consumed = (
                await _pairwise_consumed(ix, spec=spec, scenarios=selected)
                if analyze_consumed
                else []
            )

            # Replace only conflicts that touch selected scenarios (keep other conflicts intact)
            selected_names = {sc.name for sc in selected}
            kept: list[ScenarioConflict] = [
                c
                for c in spec.scenario_conflicts
                if (c.scenario_name1 not in selected_names)
                and (c.scenario_name2 not in selected_names)
            ]

            spec.scenario_conflicts = _dedup_conflicts(
                kept + list(overlaps) + list(consumed)
            )

    try:
        asyncio.run(_run())
    except RuntimeError:
        # If already in an event loop (rare for CLI), fallback.
        loop = asyncio.get_event_loop()
        loop.run_until_complete(_run())

    _refresh_scenario_stats(spec)

    if not json_only:
        n_missing = sum(isinstance(sc.component_status, Missing) for sc in selected)
        n_incon = sum(isinstance(sc.component_status, Inconsistent) for sc in selected)
        sel = {s.name for s in selected}
        n_conf = sum(
            1
            for c in spec.scenario_conflicts
            if (c.scenario_name1 in sel) or (c.scenario_name2 in sel)
        )
        console.print(
            f"[dim]Scenarios checked={len(selected)}  missing={n_missing}  inconsistent={n_incon}  conflicts={n_conf}[/dim]"
        )