#
#   Imandra Inc.
#
#   speclogician/state/change.py
#   

from __future__ import annotations

from rich.panel import Panel
from rich.text import Text
from rich.console import Group
from pydantic import BaseModel, Field, computed_field, model_validator
from typing import Literal, Optional

from ..modeling.domain import PredicateType
from ..modeling.scenario import ScenarioDelta, BoolDelta

ProcessChangeStage = Literal["apply_change", "analyze_change", "calc_diff", "ok"]

class StateChange(BaseModel):
    """ StateChange """

    meta : str | None = None

    @computed_field
    @property
    def kind(self) -> str:
        return self.__class__.__name__

    def short_str(self) -> str:
        return f"[bold deep_blue_sky]{self.kind}[/bold deep_blue_sky]"

    def __rich__(self) -> Panel:
        rows: list[Text] = []
        for k, v in self.model_dump(exclude_none=True, exclude={"kind"}).items():
            rows.append(Text(f"{k} = {v!r}"))
        return Panel(
            Group(*rows),
            title=f"[italic]{self.__class__.__name__}[/]",
            title_align="left",
        )

# ----------------------------
# Domain Base
# ----------------------------

class DomainModelBaseEdit(StateChange):
    """ Edit the base component of domain model """
    new_base_src : str = Field(..., description="IML source code of the base component")

    def short_str(self) -> str:
        return "[bold deep_sky_blue]DomainModelBaseEdit[/bold deep_sky_blue]"

    def __str__ (self) -> str:
        return ""

# ----------------------------
# Predicates
# ----------------------------

class PredicateAdd(StateChange):
    """ Add a new predicate """
    pred_type : PredicateType = Field(..., description="Predicate add")
    pred_name : str = Field(..., description="predicate name (no spaces). no signature.")
    pred_src : str = Field(..., description="predicate source code")

    def short_str(self):
        """ """
        return "[bold deep_blue_sky]PredicateAdd[/bold deep_blue_sky]"

    def __str__(self):
        return ""

class PredicateEdit(StateChange):
    """ Edit an existing predicate """
    pred_name : str = Field(..., description="Predicate name (no spaces)")
    pred_src : str = Field(..., description="Predicate source code. No signature.")

    def short_str(self):
        return "[bold deep_blue_sky]PredicateEdit[/bold deep_blue_sky]: [bold]name[/bold]={self.pred_name}"

    def __str__(self):
        return f'PredicateEdit: {self.pred_name}'
    
class PredicateRemove(StateChange):
    """ Remove an existing predicate """
    pred_name : str = Field(..., description="Predicate name to remove")

    def short_str(self) -> str:
        return f'PredicateRemove: name = {self.pred_name}'

    def __str__(self) -> str:
        return 'PredicateRemove'

# ----------------------------
# Transitions
# ----------------------------

class TransitionAdd(StateChange):
    """ Add a new transition function """
    trans_name : str = Field(..., description="Transition name")
    trans_src : str = Field(..., description="Transition source code. No signature.")

    def short_str(self) -> str:
        return ''

    def __str__(self) -> str:
        return 'TransitionAdd: '
    
class TransitionEdit(StateChange):
    """ Edit an existing transition function """
    trans_name : str = Field(..., description="Transition name")
    trans_src : str = Field(..., description="Updated transition code. No signature.")

    def short_str(self):
        return f"[bold deep_blue_sky]TransitionEdit[/bold deep_blue_sky]: [bold]trans_name[/bold]: {self.trans_name}"

    def __str__(self):
        return "TransitionEdit: "

class TransitionRemove(StateChange):
    """ Remove an existing transition function """
    trans_name : str = Field(..., description="Transition name")

    def short_str(self):
        return f"[bold deep_blue_sky]TransitionRemove[/bold deep_blue_sky]: [bold]trans_name[/bold]: {self.trans_name}"

    def __str__(self):
        return "TransitionRemove: "

# ----------------------------
# Scenarios
# ----------------------------

class ScenarioAdd(StateChange):
    """
    Add a new scenario using set-delta semantics.
    For creation, deltas must be add-only (remove must be empty).
    """
    scenario_name: str
    given: Optional[ScenarioDelta] = None
    when: Optional[ScenarioDelta] = None
    then: Optional[ScenarioDelta] = None

    is_initial: BoolDelta | None = None

    @model_validator(mode="after")
    def _add_only(self) -> "ScenarioAdd":
        for sec_name, delta in (("given", self.given), ("when", self.when), ("then", self.then)):
            if delta is not None and not delta.is_add_only():
                raise ValueError(
                    f"ScenarioAdd.{sec_name}: 'remove' is not allowed when creating a scenario"
                )
        # If provided, require it to actually set something (avoid no-op payloads)
        if self.is_initial is not None and self.is_initial.is_empty():
            raise ValueError("ScenarioAdd.is_initial has no effect: set must be true/false or omit the field.")
        return self

    def short_str(self) -> str:
        init_s = ""
        if self.is_initial is not None and self.is_initial.set is not None:
            init_s = f", is_initial={self.is_initial.set}"
        return f"ScenarioAdd: name={self.scenario_name}{init_s}"



Section = Literal["given", "when", "then"]

class ScenarioEdit(StateChange):
    scenario_name: str
    given: ScenarioDelta | None = None
    when: ScenarioDelta | None = None
    then: ScenarioDelta | None = None

    # NEW: toggle/set initialness
    is_initial: BoolDelta | None = None

    @model_validator(mode="after")
    def _must_have_some_delta(self) -> "ScenarioEdit":
        deltas = [d for d in (self.given, self.when, self.then) if d is not None]
        bool_deltas = [d for d in (self.is_initial,) if d is not None]

        if not deltas and not bool_deltas:
            raise ValueError("At least one of given/when/then/is_initial must be provided.")

        # detect all-noop
        list_noop = (not deltas) or all(d.is_empty() for d in deltas)
        bool_noop = (not bool_deltas) or all(d.is_empty() for d in bool_deltas)

        if list_noop and bool_noop:
            raise ValueError("ScenarioEdit has no effect: all provided deltas are empty.")

        return self

    def short_str(self) -> str:
        parts: list[str] = []
        for sec in ("given", "when", "then"):
            delta = getattr(self, sec)
            if delta:
                if delta.add:
                    parts.append(f"{sec} +{delta.add}")
                if delta.remove:
                    parts.append(f"{sec} -{delta.remove}")

        if self.is_initial is not None and self.is_initial.set is not None:
            parts.append(f"is_initial -> {self.is_initial.set}")

        return f"ScenarioStepsEdit[{self.scenario_name}]: " + ", ".join(parts)

class ScenarioRemove(StateChange):
    """ Remove an existing scenario """
    scenario_name : str = Field(..., description="Name of the scenario to remove")
    
    def short_str(self):
        return f"[bold deep_blue_sky]ScenarioRemove[/bold deep_blue_sky]: [bold]Name:[/bold]{self.scenario_name}"

# ----------------------------
# Test traces
# ----------------------------

class RemoveTraceArtifact(StateChange):
    """Remove an existing trace artifact (TestTrace or LogTrace)."""
    art_id: str = Field(description="Artifact ID")

    def short_str(self) -> str:
        return (
            f"[bold deep_blue_sky]RemoveTraceArtifact[/bold deep_blue_sky]: "
            f"[bold]ArtID[/bold]: {self.art_id}"
        )

class AddTestTrace(StateChange):
    """Add a new TestTrace."""
    art_id: str | None = Field(description="Artifact ID")

    # TraceArtifact core
    given: str = Field(description="Concrete state value (IML expression)")
    when: Optional[str] = Field(default=None, description="Optional concrete action value (IML expression)")
    then: Optional[str] = Field(default=None, description="Optional resulting state value (IML expression)")
    time: str = Field(default="", description="Optional timestamp")

    # TestTrace fields
    name: str = Field(description="Test case name (not unique)")
    filepath: str = Field(default="", description="Original test filepath")
    language: Optional[str] = Field(default=None, description="Language of the original test source")
    contents: Optional[str] = Field(default=None, description="Original test contents")

    def short_str(self) -> str:
        return (
            f"[bold deep_blue_sky]AddTestTrace[/bold deep_blue_sky]: "
            f"[bold]ArtID[/bold]: {self.art_id}"
        )

class EditTestTrace(StateChange):
    """Edit an existing TestTrace (patch semantics: only non-None fields are applied)."""
    art_id: str = Field(description="Artifact ID")

    # TraceArtifact core (patch)
    given: Optional[str] = Field(default=None, description="Concrete state value (IML expression)")
    when: Optional[str] = Field(default=None, description="Optional concrete action value (IML expression)")
    then: Optional[str] = Field(default=None, description="Optional resulting state value (IML expression)")
    time: Optional[str] = Field(default=None, description="Optional timestamp")

    # TestTrace fields (patch)
    name: Optional[str] = Field(default=None, description="Test case name (not unique)")
    filepath: Optional[str] = Field(default=None, description="Original test filepath")
    language: Optional[str] = Field(default=None, description="Language of the original test source")
    contents: Optional[str] = Field(default=None, description="Original test contents")

    @model_validator(mode="after")
    def _must_change_something(self) -> EditTestTrace:
        if (
            self.given is None
            and self.when is None
            and self.then is None
            and self.time is None
            and self.name is None
            and self.filepath is None
            and self.language is None
            and self.contents is None
        ):
            raise ValueError(
                "EditTestTrace must provide at least one of: "
                "given, when, then, time, name, filepath, language, or contents"
            )
        return self

    def short_str(self) -> str:
        return (
            f"[bold deep_blue_sky]EditTestTrace[/bold deep_blue_sky]: "
            f"[bold]ArtID[/bold]: {self.art_id}"
        )

# ----------------------------
# Log traces
# ----------------------------

class AddLogTrace(StateChange):
    """Add a new LogTrace."""
    art_id: str | None = Field(description="Artifact ID")

    # TraceArtifact core
    given: str = Field(description="Concrete state value (IML expression)")
    when: Optional[str] = Field(default=None, description="Optional concrete action value (IML expression)")
    then: Optional[str] = Field(default=None, description="Optional resulting state value (IML expression)")
    time: str = Field(default="", description="Optional timestamp")

    # LogTrace fields
    filename: str = Field(description="Log filename (not unique)")
    contents: str = Field(description="Log entry contents")

    def short_str(self) -> str:
        return (
            f"[bold deep_blue_sky]AddLogTrace[/bold deep_blue_sky]: "
            f"[bold]ArtID[/bold]: {self.art_id}"
        )

class EditLogTrace(StateChange):
    """Edit an existing LogTrace (patch semantics: only non-None fields are applied)."""
    art_id: str = Field(description="Artifact ID")

    # TraceArtifact core (patch)
    given: Optional[str] = Field(default=None, description="Concrete state value (IML expression)")
    when: Optional[str] = Field(default=None, description="Optional concrete action value (IML expression)")
    then: Optional[str] = Field(default=None, description="Optional resulting state value (IML expression)")
    time: Optional[str] = Field(default=None, description="Optional timestamp")

    # LogTrace fields (patch)
    filename: Optional[str] = Field(default=None, description="Log filename (not unique)")
    contents: Optional[str] = Field(default=None, description="Log entry contents")

    @model_validator(mode="after")
    def _must_change_something(self) -> EditLogTrace:
        if (
            self.given is None
            and self.when is None
            and self.then is None
            and self.time is None
            and self.filename is None
            and self.contents is None
        ):
            raise ValueError(
                "EditLogTrace must provide at least one of: "
                "given, when, then, time, filename, or contents"
            )
        return self

    def short_str(self) -> str:
        return (
            f"[bold deep_blue_sky]EditLogTrace[/bold deep_blue_sky]: "
            f"[bold]ArtID[/bold]: {self.art_id}"
        )

# ----------------------------
# RemoveDataArtifact
# ----------------------------

class RemoveDataArtifact(StateChange):
    """ Remove an existing data artifact """
    art_id: str = Field(description="Artifact ID")
    def short_str(self) -> str:
        return "[bold deep_blue_sky]RemoveDataArtifact[/bold deep_blue_sky]"

# ----------------------------
# Documentation references
# ----------------------------

class AddDocRef(StateChange):
    """Add a new documentation reference artifact."""
    art_id: str | None = Field(description="Artifact ID")
    text: str = Field(description="Documentation text")
    meta: str | None = Field(default=None, description="Optional metadata")

    def short_str(self) -> str:
        return f"[bold deep_blue_sky]AddDocRef[/bold deep_blue_sky]: [bold]ArtID[/bold]: {self.art_id}"

class EditDocRef(StateChange):
    """Edit an existing DocRef."""
    art_id: str = Field(description="Artifact ID")
    text: str | None = Field(default=None, description="New documentation text")
    meta: str | None = Field(default=None, description="New metadata")

    @model_validator(mode="after")
    def _must_change_something(self) -> "EditDocRef":
        if self.text is None and self.meta is None:
            raise ValueError("EditDocRef must provide at least one of: text or meta")
        return self

    def short_str(self) -> str:
        return f"[bold deep_blue_sky]EditDocRef[/bold deep_blue_sky]: [bold]ArtID[/bold]: {self.art_id}"
    
# ----------------------------
# Source code references
# ----------------------------

class AddSrcCodeRef(StateChange):
    """Add a new source code reference artifact."""
    art_id: str | None = Field(description="Artifact ID")
    src_code: str = Field(description="Source code")
    language: str | None = Field(default=None, description="Programming language")
    file_path: str | None = Field(default=None, description="Optional file path")
    iml_code: str | None = Field(default=None, description="Optional IML formalization")
    meta: str | None = Field(default=None, description="Optional metadata")

    def short_str(self) -> str:
        return f"[bold deep_blue_sky]AddSrcCodeRef[/bold deep_blue_sky]: [bold]ArtID[/bold]: {self.art_id}"
    
class EditSrcCodeRef(StateChange):
    """Edit an existing SrcCodeRef."""
    art_id: str = Field(description="Artifact ID")
    src_code: str | None = Field(default=None, description="Updated source code")
    language: str | None = Field(default=None, description="Updated language")
    file_path: str | None = Field(default=None, description="Updated file path")
    iml_code: str | None = Field(default=None, description="Updated IML formalization")
    meta: str | None = Field(default=None, description="Updated metadata")

    @model_validator(mode="after")
    def _must_change_something(self) -> EditSrcCodeRef:
        if (
            self.src_code is None
            and self.language is None
            and self.file_path is None
            and self.iml_code is None
            and self.meta is None
        ):
            raise ValueError(
                "EditSrcCodeRef must provide at least one of: "
                "src_code, language, file_path, iml_code, or meta"
            )
        return self

    def short_str(self) -> str:
        return f"[bold deep_blue_sky]EditSrcCodeRef[/bold deep_blue_sky]: [bold]ArtID[/bold]: {self.art_id}"

# ----------------------------
# Links in the ArtifactMap
# ----------------------------

class LinkArtifactsComponents(StateChange):
    """ Add a link between a non-executable arifact (e.g. source code or documentation) and a model component """
    art_id : str = Field(description="Artifact ID")
    comp_name : str = Field(description="Component name")

    def short_str(self) -> str:
        return "[bold deep_blue_sky]LinkArtifactComponents[/bold deep_blue_sky]"

class RemoveArtComponentLink(StateChange):
    """ Remove a link between a non-executable arifact (e.g. source code or documentation) and a model component """
    art_id : str = Field(description="Artifact ID")
    comp_name : str = Field(description="Component name")

    def short_str(self) -> str:
        return "[bold deep_blue_sky]RemoveArtComponentLink[/bold deep_blue_sky]"


# ---------------------------------------------------------------------------
# Batch deserialization
# ---------------------------------------------------------------------------

_CHANGE_REGISTRY: dict[str, type[StateChange]] = {
    "DomainModelBaseEdit": DomainModelBaseEdit,
    "PredicateAdd": PredicateAdd,
    "PredicateEdit": PredicateEdit,
    "PredicateRemove": PredicateRemove,
    "TransitionAdd": TransitionAdd,
    "TransitionEdit": TransitionEdit,
    "TransitionRemove": TransitionRemove,
    "ScenarioAdd": ScenarioAdd,
    "ScenarioEdit": ScenarioEdit,
    "ScenarioRemove": ScenarioRemove,
    "AddTestTrace": AddTestTrace,
    "EditTestTrace": EditTestTrace,
    "AddLogTrace": AddLogTrace,
    "EditLogTrace": EditLogTrace,
    "RemoveTraceArtifact": RemoveTraceArtifact,
    "AddDocRef": AddDocRef,
    "EditDocRef": EditDocRef,
    "AddSrcCodeRef": AddSrcCodeRef,
    "EditSrcCodeRef": EditSrcCodeRef,
    "RemoveDataArtifact": RemoveDataArtifact,
    "LinkArtifactsComponents": LinkArtifactsComponents,
    "RemoveArtComponentLink": RemoveArtComponentLink,
}


def parse_change(obj: dict) -> StateChange:
    """
    Deserialize a dict with a ``kind`` discriminator into a ``StateChange``.

    Expected format::

        {"kind": "PredicateAdd", "pred_type": "state", "pred_name": "x_gt_0", "pred_src": "s.x > 0"}
    """
    kind = obj.get("kind")
    if kind is None:
        raise ValueError("Each change object must have a 'kind' field")
    cls = _CHANGE_REGISTRY.get(kind)
    if cls is None:
        valid = ", ".join(sorted(_CHANGE_REGISTRY))
        raise ValueError(f"Unknown change kind '{kind}'. Valid kinds: {valid}")
    # Remove 'kind' before validating (it's a computed field, not an input)
    fields = {k: v for k, v in obj.items() if k != "kind"}
    return cls.model_validate(fields)