#
#   Imandra Inc.
#
#   speclogician/logic/main.py
#

from __future__ import annotations

from ..modeling.spec import Spec
from ..data.container import ArtifactContainer
from ..data.mapping import ArtifactMap

from ..state.change import (
    StateChange,
    DomainModelBaseEdit,
    PredicateAdd,
    PredicateEdit,
    PredicateRemove,
    TransitionAdd,
    TransitionEdit,
    TransitionRemove,
    ScenarioAdd,
    ScenarioEdit,
    ScenarioRemove,
    AddTestTrace,
    EditTestTrace,
    AddLogTrace,
    EditLogTrace
)

from .lib.domain_model import check_domain_model
from .lib.predicates import check_predicates
from .lib.transitions import check_transitions
from .lib.scenarios import check_scenarios
from .lib.traces import check_traces
from .lib.complement import maybe_update_spec_complement
from .lib import update_reachability_graph

from ..utils import console


def _hdr(title: str, subtitle: str | None = None, json_only: bool = False) -> None:
    if json_only:
        return
    from rich.panel import Panel
    from rich.text import Text

    console.print(
        Panel(
            Text(subtitle or ""),
            title=f"[bold]Logician[/bold] · {title}",
            border_style="cyan",
        )
    )


def _note(msg: str, json_only: bool = False) -> None:
    if not json_only:
        console.print(msg)

def analyze_change(
    spec: Spec,
    art_cont: ArtifactContainer,
    art_map: ArtifactMap,
    change: StateChange,
    json_only: bool = False,
) -> None:
    """
    Top-level entry for logical analysis post-change event.

    Only performs logic checks when relevant.
    All other change types are treated as no-ops here.
    """

    # ----------------------------
    # Domain base model changed
    # ----------------------------
    if isinstance(change, DomainModelBaseEdit):
        _hdr("DomainModelBaseEdit", "Rechecking domain model and all dependent logic", json_only=json_only)
        check_domain_model(spec, json_only=json_only)
        check_predicates(spec, json_only=json_only)
        check_transitions(spec, json_only=json_only)
        check_scenarios(spec, json_only=json_only)
        check_traces(spec, art_cont=art_cont, art_map=art_map, traces=[], json_only=json_only)

        try:
            maybe_update_spec_complement(spec)
        except Exception as e:
            _note(f"[yellow]⚠[/yellow] Complement update failed (non-fatal): {type(e).__name__}: {e}", json_only=json_only)

        try:
            update_reachability_graph(spec, json_only=json_only)
        except Exception as e:
            _note(f"[yellow]⚠[/yellow] Reachability update failed (non-fatal): {type(e).__name__}: {e}", json_only=json_only)

        _note("[green]✓[/green] Domain analysis complete", json_only=json_only)
        return

    # ----------------------------
    # Predicates changed
    # ----------------------------
    if isinstance(change, (PredicateAdd, PredicateEdit)):
        _hdr("Predicate change", f"Predicate: [bold]{change.pred_name}[/bold]", json_only=json_only)
        check_predicates(spec, names=[change.pred_name], json_only=json_only)
        check_scenarios(spec, predicates=[change.pred_name], json_only=json_only)
        check_traces(spec, art_cont=art_cont, art_map=art_map, traces=[], json_only=json_only)

        try:
            maybe_update_spec_complement(spec)
        except Exception as e:
            _note(f"[yellow]⚠[/yellow] Complement update failed (non-fatal): {type(e).__name__}: {e}", json_only=json_only)

        try:
            update_reachability_graph(spec, json_only=json_only)
        except Exception as e:
            _note(f"[yellow]⚠[/yellow] Reachability update failed (non-fatal): {type(e).__name__}: {e}", json_only=json_only)


        _note("[green]✓[/green] Predicate analysis complete", json_only=json_only)
        return

    if isinstance(change, PredicateRemove):
        _hdr("Predicate removed", f"[bold]{change.pred_name}[/bold]", json_only=json_only)
        check_scenarios(spec, predicates=[change.pred_name], json_only=json_only)
        check_traces(spec, art_cont=art_cont, art_map=art_map, traces=[], json_only=json_only)

        try:
            maybe_update_spec_complement(spec)
        except Exception as e:
            _note(f"[yellow]⚠[/yellow] Complement update failed (non-fatal): {type(e).__name__}: {e}", json_only=json_only)

        try:
            update_reachability_graph(spec, json_only=json_only)
        except Exception as e:
            _note(f"[yellow]⚠[/yellow] Reachability update failed (non-fatal): {type(e).__name__}: {e}", json_only=json_only)

        _note("[green]✓[/green] Scenario refresh complete", json_only=json_only)
        return

    # ----------------------------
    # Transitions changed
    # ----------------------------
    if isinstance(change, (TransitionAdd, TransitionEdit)):
        _hdr("Transition change", f"Transition: [bold]{change.trans_name}[/bold]", json_only=json_only)
        check_transitions(spec, names=[change.trans_name], json_only=json_only)
        check_scenarios(spec, transitions=[change.trans_name], json_only=json_only)
        check_traces(spec, art_cont=art_cont, art_map=art_map, traces=[], json_only=json_only)

        try:
            maybe_update_spec_complement(spec)
        except Exception as e:
            _note(f"[yellow]⚠[/yellow] Complement update failed (non-fatal): {type(e).__name__}: {e}", json_only=json_only)

        try:
            update_reachability_graph(spec, json_only=json_only)
        except Exception as e:
            _note(f"[yellow]⚠[/yellow] Reachability update failed (non-fatal): {type(e).__name__}: {e}", json_only=json_only)

        _note("[green]✓[/green] Transition analysis complete", json_only=json_only)
        return

    if isinstance(change, TransitionRemove):
        _hdr("Transition removed", f"[bold]{change.trans_name}[/bold]", json_only=json_only)
        check_scenarios(spec, transitions=[change.trans_name], json_only=json_only)
        check_traces(spec, art_cont=art_cont, art_map=art_map, traces=[], json_only=json_only)

        try:
            maybe_update_spec_complement(spec)
        except Exception as e:
            _note(f"[yellow]⚠[/yellow] Complement update failed (non-fatal): {type(e).__name__}: {e}", json_only=json_only)

        try:
            update_reachability_graph(spec, json_only=json_only)
        except Exception as e:
            _note(f"[yellow]⚠[/yellow] Reachability update failed (non-fatal): {type(e).__name__}: {e}", json_only=json_only)

        _note("[green]✓[/green] Scenario refresh complete", json_only=json_only)
        return

    # ----------------------------
    # Scenarios changed
    # ----------------------------
    if isinstance(change, (ScenarioAdd, ScenarioEdit)):
        _hdr("Scenario change", f"Scenario: [bold]{change.scenario_name}[/bold]", json_only=json_only)
        check_scenarios(spec, names=[change.scenario_name], json_only=json_only)
        check_traces(spec, art_cont=art_cont, art_map=art_map, traces=[], json_only=json_only)

        try:
            maybe_update_spec_complement(spec)
        except Exception as e:
            _note(f"[yellow]⚠[/yellow] Complement update failed (non-fatal): {type(e).__name__}: {e}", json_only=json_only)

        try:
            update_reachability_graph(spec, json_only=json_only)
        except Exception as e:
            _note(f"[yellow]⚠[/yellow] Reachability update failed (non-fatal): {type(e).__name__}: {e}", json_only=json_only)


        _note("[green]✓[/green] Scenario analysis complete", json_only=json_only)
        return

    if isinstance(change, ScenarioRemove):
        _hdr("Scenario removed", "Refreshing scenarios, predicates, and traces", json_only=json_only)
        check_scenarios(spec, json_only=json_only)
        check_predicates(spec, json_only=json_only)
        check_traces(spec=spec, art_cont=art_cont, art_map=art_map, traces=[], json_only=json_only)

        try:
            maybe_update_spec_complement(spec)
        except Exception as e:
            _note(f"[yellow]⚠[/yellow] Complement update failed (non-fatal): {type(e).__name__}: {e}", json_only=json_only)

        try:
            update_reachability_graph(spec, json_only=json_only)
        except Exception as e:
            _note(f"[yellow]⚠[/yellow] Reachability update failed (non-fatal): {type(e).__name__}: {e}", json_only=json_only)


        _note("[green]✓[/green] Full refresh complete", json_only=json_only)
        return

    # ----------------------------
    # Trace artifacts changed
    # ----------------------------
    if isinstance(change, AddTestTrace):
        if not art_cont.test_traces:
            return

        tt = art_cont.test_traces[-1]
        # best-effort sanity check; no-op if mismatch
        if change.art_id is not None and tt.art_id != change.art_id:
            pass

        assert tt.art_id is not None

        check_traces(
            spec=spec,
            art_cont=art_cont,
            art_map=art_map,
            traces=[tt.art_id],
            json_only=json_only,
        )

        _note("[green]✓[/green] Trace analysis complete", json_only=json_only)
        return

    if isinstance(change, AddLogTrace):
        if not art_cont.log_traces:
            return

        lt = art_cont.log_traces[-1]
        if change.art_id is not None and lt.art_id != change.art_id:
            pass

        assert lt.art_id is not None

        check_traces(
            spec=spec,
            art_cont=art_cont,
            art_map=art_map,
            traces=[lt.art_id],
            json_only=json_only,
        )

        _note("[green]✓[/green] Trace analysis complete", json_only=json_only)
        return

    if isinstance(change, (EditTestTrace, EditLogTrace)):
        check_traces(
            spec=spec,
            art_cont=art_cont,
            art_map=art_map,
            traces=[change.art_id],
            json_only=json_only,
        )

        _note("[green]✓[/green] Trace analysis complete", json_only=json_only)
        return

    # ----------------------------
    # Explicit no-op branch
    # ----------------------------
    # All remaining change types (data artifacts, mapping edits, etc.)
    # are intentionally ignored here and handled elsewhere.
    return


def analyze_batch(
    spec: Spec,
    art_cont: ArtifactContainer,
    art_map: ArtifactMap,
    changes: list[StateChange],
    json_only: bool = False,
) -> None:
    """
    Batch-aware analysis: compute the union scope of all changes,
    then run each check function once with the merged scope.
    """

    # If any change is a DomainModelBaseEdit, do a full recheck
    if any(isinstance(c, DomainModelBaseEdit) for c in changes):
        _hdr("Batch (full recheck)", "Batch contains DomainModelBaseEdit — rechecking everything", json_only=json_only)
        check_domain_model(spec, json_only=json_only)
        check_predicates(spec, json_only=json_only)
        check_transitions(spec, json_only=json_only)
        check_scenarios(spec, json_only=json_only)
        check_traces(spec, art_cont=art_cont, art_map=art_map, traces=[], json_only=json_only)
        _try_complement_and_reachability(spec, json_only=json_only)
        _note("[green]✓[/green] Batch analysis complete (full recheck)", json_only=json_only)
        return

    # Collect affected names/ids across all changes
    pred_names: list[str] = []
    trans_names: list[str] = []
    scenario_names: list[str] = []
    trace_ids: list[str] = []
    full_scenario_refresh = False
    full_predicate_refresh = False

    for c in changes:
        if isinstance(c, (PredicateAdd, PredicateEdit, PredicateRemove)):
            pred_names.append(c.pred_name)
        if isinstance(c, (TransitionAdd, TransitionEdit, TransitionRemove)):
            trans_names.append(c.trans_name)
        if isinstance(c, (ScenarioAdd, ScenarioEdit)):
            scenario_names.append(c.scenario_name)
        if isinstance(c, ScenarioRemove):
            full_scenario_refresh = True
            full_predicate_refresh = True
        if isinstance(c, (AddTestTrace, EditTestTrace)):
            if hasattr(c, "art_id") and c.art_id is not None:
                trace_ids.append(c.art_id)
        if isinstance(c, (AddLogTrace, EditLogTrace)):
            if hasattr(c, "art_id") and c.art_id is not None:
                trace_ids.append(c.art_id)

    has_logic_changes = pred_names or trans_names or scenario_names or full_scenario_refresh

    if not has_logic_changes and not trace_ids:
        _note("[dim]Batch: no logic-relevant changes, skipping analysis[/dim]", json_only=json_only)
        return

    _hdr("Batch analysis", f"{len(changes)} changes", json_only=json_only)

    # Predicates
    if pred_names or full_predicate_refresh:
        if full_predicate_refresh:
            check_predicates(spec, json_only=json_only)
        else:
            check_predicates(spec, names=pred_names, json_only=json_only)

    # Transitions
    if trans_names:
        check_transitions(spec, names=trans_names, json_only=json_only)

    # Scenarios — affected by pred/trans/scenario changes
    if has_logic_changes:
        if full_scenario_refresh:
            check_scenarios(spec, json_only=json_only)
        else:
            check_scenarios(
                spec,
                names=scenario_names,
                predicates=pred_names,
                transitions=trans_names,
                json_only=json_only,
            )

    # Traces
    if trace_ids or has_logic_changes:
        check_traces(
            spec=spec,
            art_cont=art_cont,
            art_map=art_map,
            traces=trace_ids,
            json_only=json_only,
        )

    # Complement + reachability (once)
    if has_logic_changes:
        _try_complement_and_reachability(spec, json_only=json_only)

    _note("[green]✓[/green] Batch analysis complete", json_only=json_only)


def _try_complement_and_reachability(spec: Spec, json_only: bool = False) -> None:
    """Helper: run complement + reachability, swallowing non-fatal errors."""
    try:
        maybe_update_spec_complement(spec)
    except Exception as e:
        _note(f"[yellow]⚠[/yellow] Complement update failed (non-fatal): {type(e).__name__}: {e}", json_only=json_only)

    try:
        update_reachability_graph(spec, json_only=json_only)
    except Exception as e:
        _note(f"[yellow]⚠[/yellow] Reachability update failed (non-fatal): {type(e).__name__}: {e}", json_only=json_only)