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

from __future__ import annotations

from typing import TypeAlias, Sequence
from pydantic import BaseModel, Field, field_validator, model_validator

from .component import ModelComponent, validate_iml_identifier


class BoolDelta(BaseModel):
    # None means "no change"
    set: bool | None = None

    def is_empty(self) -> bool:
        return self.set is None


class ScenarioDelta(BaseModel):
    """We'll use this to edit Scenarios."""
    add: list[str] = Field(default_factory=list)
    remove: list[str] = Field(default_factory=list)

    is_initial: BoolDelta | None = None

    @field_validator("add", "remove")
    @classmethod
    def _dedup_and_strip(cls, xs: list[str]) -> list[str]:
        # normalize, drop empties, dedup preserving order
        out: list[str] = []
        seen: set[str] = set()
        for x in xs:
            x = x.strip()
            if not x or x in seen:
                continue
            out.append(x)
            seen.add(x)
        return out

    @field_validator("add", "remove")
    @classmethod
    def _validate_identifiers(cls, xs: list[str]) -> list[str]:
        # Reuse the canonical IML identifier validation
        for x in xs:
            validate_iml_identifier(x)
        return xs

    @model_validator(mode="after")
    def _no_overlap(self) -> "ScenarioDelta":
        overlap = set(self.add) & set(self.remove)
        if overlap:
            raise ValueError(f"Names cannot be both added and removed: {sorted(overlap)}")
        return self

    def is_empty(self) -> bool:
        return not self.add and not self.remove

    def is_add_only(self) -> bool:
        return not self.remove

class Valid(BaseModel):
    """ The scenario's valid """
    pass

class Inconsistent(BaseModel):
    """ Scenario is inconsistent - the combination of predicates is unsatisfiable """
    pass

class Missing(BaseModel):
    """ Referenced model components are missing """
    missing_preds : list[str] = Field(default_factory=list)
    missing_trans : list[str] = Field(default_factory=list)

ComponentStatus : TypeAlias = Valid | Inconsistent | Missing

class Scenario(ModelComponent):
    """ 
    Scenario contains 
    """

    name : str

    is_initial: bool = False # Will be used in reachability graph

    # Predicates and transition functions are components
    component_status : ComponentStatus = Valid()

    given: Sequence[str] = Field(default_factory=list)
    when: Sequence[str] = Field(default_factory=list)
    then: Sequence[str] = Field(default_factory=list)


    # When we check the DomainModel and this is missing, 
    # we'll then set these fields
    preds_missing : list[str] = Field(default_factory=list)
    trans_missing : list[str] = Field(default_factory=list)

    # These contain consistency checks
    # why precisely they're inconsistent
    given_preds_consistent : bool = True
    when_preds_consistent : bool = True
    all_preds_consistent : bool = True

    def pred_calls_to_iml (self) -> str:
        """
        return &&'ed string of all predicates
        """

        given_str = "\n&& ".join(map(lambda x: f"{x} s", self.given))
        when_str = "\n&& ".join(map(lambda x: f"{x} s a", self.when))

        match (len(self.given), len(self.when)):
            case (0, 0):
                return ""
            case (_, 0):
                return given_str
            case (0, _):
                return when_str
            case (_, _):
                return f"{given_str} \n&& {when_str}"

    def eval_transition (self, state_val:str, act_val:str) -> str:
        """
        Given concrete state and action values, evaluate the set of transitions
        """

        if len(self.then) == 0:
            return state_val

        trans_iml = ""
        for t in self.then:
            trans_iml += f"\nlet s = {t} s a in"

        trans_iml += "\n"

        return ''

    def preds_to_iml (self):
        """
        Return IML code of just the predicates (both types) encoded in IML
        """
        given_str = "&& \n".join(map(str, self.given))
        when_str = "&& \n".join(map(str, self.when))

        s = f"""
{given_str}

{when_str}
"""     
        return s

    def full_model_to_iml(self):
        """
        Create the full model (we omit the 'else' at the end because this is intended to be used with 
        other scenarios
        """
        given_str = "&& \n".join(map(str, self.given))
        when_str = "&& \n".join(map(str, self.when))
        then_str = "\n".join(map(str, self.then))

        s = f"""
(* Scenario: {self.name} *)
let given_true = 
{given_str} 
in let when_true = 
{when_str}
in if given_true && when_true then
({then_str})
"""
        return s