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

from __future__ import annotations

import asyncio

from speclogician.logic.api import SpecLogicianImandraX
from speclogician.modeling.spec import Spec
from speclogician.modeling.scenario import Scenario, Missing
from speclogician.modeling.reachability import (
    ScenarioReachabilityEdge,
    ScenarioReachabilityGraph,
    build_reachability_graph,
)

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)


def compute_scenario_reachability_graph(
    spec: Spec,
    *,
    initial_names: list[str] | None = None,
    require_to_when: bool = False,
    json_only: bool = False,
    store_on_spec: bool = True,
    witness: bool = True,
) -> ScenarioReachabilityGraph:
    """
    Compute the spec-level scenario reachability graph by:
      1) calling ImandraX to compute pairwise reachability edges
      2) building a deterministic ScenarioReachabilityGraph (succ/pred + reachable/unreachable)

    If store_on_spec=True, sets spec.scenario_reachability.

    NOTE:
      - Only "checkable" scenarios participate (not Missing, and all_preds_consistent=True).
      - 'initial_names' overrides spec.initial_scenarios (if present).
      - If no initials are provided, reachable/unreachable will be empty lists (by design).
    """
    scenarios = list(getattr(spec, "scenarios", []) or [])
    checkable = [sc for sc in scenarios if _is_checkable(sc)]

    # Node universe (checkable only)
    names = sorted({sc.name for sc in checkable if getattr(sc, "name", None)})

    by_name: dict[str, Scenario] = {sc.name: sc for sc in checkable if sc.name in names}

    edges: list[ScenarioReachabilityEdge] = []

    async def _run() -> None:
        async with SpecLogicianImandraX() as ix:
            dm = spec.domain_model

            # O(n^2) pairwise checks; OK for now (optimize later with batching/heuristics)
            for src in names:
                for dst in names:
                    if src == dst:
                        continue

                    r = await ix.check_scenarios_reachable(
                        domain_model=dm,
                        sc_from=by_name[src],
                        sc_to=by_name[dst],
                        require_to_when=require_to_when,
                    )

                    if r.is_reachable:
                        edges.append(
                            ScenarioReachabilityEdge(
                                src=src,
                                dst=dst,
                                witness=(r.model_src if witness else None),
                            )
                        )

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

    # Initial scenarios: prefer explicit argument, else spec.initial_scenarios if you have it
    initials = list(
        initial_names
        if initial_names is not None
        else (getattr(spec, "initial_scenarios", []) or [])
    )

    g = build_reachability_graph(
        scenario_names=names,
        edges=edges,
        initials=initials,
    )

    if store_on_spec:
        spec.scenario_reachability = g

    if not json_only:
        # Keep this lightweight; caller can do richer printing
        # (Avoid importing console here unless you want the UX.)
        pass

    return g

def update_reachability_graph(
    spec: Spec,
    json_only: bool = False
):
    """
    Convenience wrapper: compute + store the scenario reachability graph on `spec`
    using whatever defaults are present on the Spec (e.g. initial_scenarios),
    and return the computed graph.

    This is intended to be called after changes / analysis passes that may have
    modified scenarios, predicates, transitions, or consistency flags.
    """

    # Allow the Spec to carry default reach semantics (optional, future-proof).
    # If absent, we fall back to the "standard" defaults.
    require_to_when = bool(getattr(spec, "reach_require_to_when", False))
    store_witness = bool(getattr(spec, "reach_store_witness", True))

    # Compute and store on spec
    spec.scenario_reachability = compute_scenario_reachability_graph(
        spec,
        initial_names=spec.initial_scenarios(),
        require_to_when=require_to_when,
        json_only=json_only,
        store_on_spec=True,
        witness=store_witness,
    )
