#
#   Imandra Inc.
#
#   speclogician/modeling/spec.py
#

from typing import Optional
from pydantic import BaseModel, Field
from collections import defaultdict

from .domain import DomainModel
from .complement import ScenarioComplement
from .scenario import Scenario, ScenarioDelta, BoolDelta
from .conflict import ScenarioConflict
from .reachability import ScenarioReachabilityGraph


def _scenario_guard_iml(scn) -> str:
    """
    given predicates are assumed: pred(s)
    when predicates are assumed: pred(s)(a) or pred(s, a) style; we generate pred s a.
    """
    parts: list[str] = []

    for p in (scn.given or []):
        p = (p or "").strip()
        if p:
            parts.append(f"{p} s")

    for p in (scn.when or []):
        p = (p or "").strip()
        if p:
            parts.append(f"{p} s a")

    return " && ".join(parts) if parts else "true"


def _scenario_then_iml(scn) -> str:
    """
    then is a list of transition function names.
    Each transition is assumed: trans(s)(a) -> state; we generate trans s a.
    Applied sequentially.
    """
    thens = [(t or "").strip() for t in (scn.then or []) if (t or "").strip()]
    if not thens:
        return "s"

    # Apply transitions in order:
    # let s1 = t1 s a in let s2 = t2 s1 a in s2
    lines: list[str] = []
    cur = "s"
    for i, tname in enumerate(thens, start=1):
        nxt = f"s{i}"
        lines.append(f"let {nxt} = {tname} {cur} a in")
        cur = nxt
    lines.append(cur)
    return "\n  ".join(lines)


class Spec(BaseModel):
    """ Spec is comprised of the Domain Model and a list of Scenarios """

    # Domain model object also has all the relevant stats on the model
    domain_model : DomainModel = Field(default_factory=DomainModel)
    
    # The list of scenarios
    scenarios : list[Scenario] = Field(default_factory=list[Scenario])

    # Are there conflicts between?
    scenario_conflicts : list[ScenarioConflict] = Field(default_factory=list[ScenarioConflict])

    # Scenario complement
    scenario_comp : ScenarioComplement | None = None

    # Scenario stats
    num_sc_total : int = 0 # Total number of scenarios
    num_sc_missing : int = 0 # Missing predicates/transitions
    num_sc_matched : int = 0 # Total number of scenarios matched vs artifacts
    num_sc_inconsistent : int = 0 # Total number of scenarios inconsistent

    # Conflicts (spec-level)
    num_sc_conflicted: int = 0   # total conflicts (all kinds)
    num_sc_overlap: int = 0      # overlap conflicts
    num_sc_consumed: int = 0     # consumed conflicts

    # Scenario reachability graph
    scenario_reachability: ScenarioReachabilityGraph | None = None

    def conflicts_index(self) -> dict[str, list[ScenarioConflict]]:
        idx: dict[str, list[ScenarioConflict]] = defaultdict(list)
        for c in self.scenario_conflicts:
            idx[c.scenario_name1].append(c)
            idx[c.scenario_name2].append(c)
        return dict(idx)

    def conflicts_for(self, scenario_name: str) -> list[ScenarioConflict]:
        return [
            c for c in self.scenario_conflicts
            if c.scenario_name1 == scenario_name or c.scenario_name2 == scenario_name
        ]

    def scenario_exists (self, name:str) -> bool:
        """
        Return True/False if scenario already exists
        """
        return any(s.name == name for s in self.scenarios)
    
    def find_scenario(self, name: str) -> Optional[Scenario]:
        """
        Return the Scenario with the given name, or None if not found.
        """
        for s in self.scenarios:
            if s.name == name:
                return s
        return None

    def initial_scenarios(self) -> list[str]:
        """
        Return list of scenarios marked 'is_initial'.
        """
        return [sc.name for sc in self.scenarios if sc.is_initial]

    def _apply_delta_strict(
        self,
        current: list[str],
        *,
        delta: ScenarioDelta,
        allowed_add: set[str],
        section_label: str,
    ) -> list[str]:
        cur = set(current)

        # --- validate removals: must already be present ---
        missing_removals = set(delta.remove) - cur
        if missing_removals:
            raise ValueError(
                f"{section_label}: cannot remove names that are not present: {sorted(missing_removals)}"
            )

        # --- validate additions only ---
        bad_adds = set(delta.add) - allowed_add
        if bad_adds:
            raise ValueError(
                f"{section_label}: cannot add unknown names: {sorted(bad_adds)}"
            )

        # optional: forbid adding something already present
        already_present = set(delta.add) & cur
        if already_present:
            raise ValueError(
                f"{section_label}: names already present: {sorted(already_present)}"
            )

        # --- apply ---
        cur.difference_update(delta.remove)
        cur.update(delta.add)

        return sorted(cur)  # deterministic
    

    def add_scenario(
        self,
        name: str,
        given: ScenarioDelta | None,
        when: ScenarioDelta | None,
        then: ScenarioDelta | None,
        is_initial: BoolDelta | None,
    ) -> None:
        """Add a new scenario. Deltas must be add-only."""
        if self.scenario_exists(name):
            raise ValueError(f"Scenario with this name already exists: [{name}]")

        allowed_given = set(self.domain_model.state_pred_names())
        allowed_when = set(self.domain_model.action_pred_names())
        allowed_then = set(self.domain_model.transition_names())

        def build_section(delta: ScenarioDelta | None, allowed: set[str], label: str) -> list[str]:
            if delta is None:
                return []
            if not delta.is_add_only():
                raise ValueError(f"{label}: 'remove' is not allowed when creating a scenario")

            bad = set(delta.add) - allowed
            if bad:
                raise ValueError(f"{label}: cannot add unknown names: {sorted(bad)}")

            return sorted(set(delta.add))  # set semantics, deterministic

        new_given = build_section(given, allowed_given, f"Scenario '{name}' given")
        new_when = build_section(when, allowed_when, f"Scenario '{name}' when")
        new_then = build_section(then, allowed_then, f"Scenario '{name}' then")

        # ---- NEW: initial flag plumbing ----
        if is_initial is not None and is_initial.set is None:
            raise ValueError("ScenarioAdd.is_initial has no effect: set must be true/false or omit the field.")

        new_is_initial = bool(is_initial.set) if (is_initial is not None and is_initial.set is not None) else False

        self.scenarios.append(
            Scenario(
                name=name,
                given=new_given,
                when=new_when,
                then=new_then,
                is_initial=new_is_initial,
            )
        )

    def edit_scenario(
        self,
        name: str,
        given: Optional[ScenarioDelta],
        when: Optional[ScenarioDelta],
        then: Optional[ScenarioDelta],
        is_initial: BoolDelta | None,
    ) -> None:
        """Edit an existing scenario with strict validation (checks only on add)."""
        scenario = self.find_scenario(name)
        if scenario is None:
            raise ValueError(f"Scenario '{name}' not found")

        allowed_given = set(self.domain_model.state_pred_names())
        allowed_when = set(self.domain_model.action_pred_names())
        allowed_then = set(self.domain_model.transition_names())

        if given is not None:
            scenario.given = self._apply_delta_strict(
                list(scenario.given),
                delta=given,
                allowed_add=allowed_given,
                section_label=f"Scenario '{name}' given",
            )

        if when is not None:
            scenario.when = self._apply_delta_strict(
                list(scenario.when),
                delta=when,
                allowed_add=allowed_when,
                section_label=f"Scenario '{name}' when",
            )

        if then is not None:
            scenario.then = self._apply_delta_strict(
                list(scenario.then),
                delta=then,
                allowed_add=allowed_then,
                section_label=f"Scenario '{name}' then",
            )

        if is_initial is not None:
            if is_initial.set is None:
                raise ValueError("ScenarioEdit.is_initial has no effect: set must be true/false or omit the field.")
            scenario.is_initial = bool(is_initial.set)
        
    def rem_scenario(self, name: str) -> None:
        """Remove a scenario by name."""
        for i, s in enumerate(self.scenarios):
            if s.name == name:
                del self.scenarios[i]
                return
        raise ValueError(f"Scenario '{name}' not found")

    def scenario_model(self):
        """
        Synthesize a model from the scenarios
        """

        scenario_logic = ""
        for s in self.scenarios:
            scenario_logic += f"{s.full_model_to_iml()}\n"

        model = f"""
let scenarios (s : state) (a : action) = 
{scenario_logic} else
false
"""
        return model
    
    def _scenarios_main_iml(self) -> str:
        """
        Compile scenarios into a single dispatcher function `main`.
        Scenarios are tried in list order (priority order).
        """
        if not self.scenarios:
            return "let main (s: state) (_a: action) : state = s"

        lines: list[str] = []
        lines.append("let main (s: state) (a: action) : state =")

        for idx, scn in enumerate(self.scenarios):
            guard = _scenario_guard_iml(scn)
            body = _scenario_then_iml(scn)

            kw = "if" if idx == 0 else "else if"
            lines.append(f"  {kw} {guard} then")
            # indent body nicely
            body_lines = body.splitlines() or ["s"]
            lines.append("    " + "\n    ".join(body_lines))

        lines.append("  else s")
        return "\n".join(lines)

    def full_model(self) -> str:
        """
        Return the full model (base + predicates + transitions + scenarios)
        as a single executable IML program.
        """
        dm = self.domain_model

        parts: list[str] = []

        # 1. Base domain
        if dm.base:
            parts.append(dm.base.strip())

        # 2. Predicates
        if dm.state_preds or dm.action_preds:
            parts.append("\n(* Predicates *)\n")
            for p in dm.state_preds:
                parts.append(p.to_iml().strip())
            for p in dm.action_preds:
                parts.append(p.to_iml().strip())

        # 3. Transitions
        if dm.transitions:
            parts.append("\n(* Transitions *)\n")
            for t in dm.transitions:
                parts.append(t.to_iml().strip())

        # 4. Scenarios
        if self.scenarios:
            parts.append("\n(* Scenario dispatcher *)\n")
            parts.append(self._scenarios_main_iml().strip())

        # Join with clean spacing
        return "\n\n".join(p for p in parts if p.strip())
