#
#   Imandra Inc.
#
#   speclogician/data/trace.py
#

from __future__ import annotations

from typing import Any
from pydantic import BaseModel, Field, model_validator

from .artifact import Artifact
from ..utils import IMLValidity


# -----------------------------------------------------------------------------
# Analysis / results
# -----------------------------------------------------------------------------

class TraceIMLValidity(BaseModel):
    iml_valid: IMLValidity
    err: str | None = None


class TraceScenarioMatchResult(BaseModel):
    scenario_name: str
    iml_valid: bool
    err: str | None = None
    when_matches: bool = False
    given_matches: bool = False
    transition_matches: bool = False


# -----------------------------------------------------------------------------
# Trace data model
# -----------------------------------------------------------------------------

class TraceStep(BaseModel):
    """
    One observed step.

    Convention:
      - If you know a pre-state and action: set state + action.
      - If you only know the action/event: set action only.
      - If you only know state snapshots: set state only.
      - If you know a post-state, represent it as the NEXT step's state.
    """
    state: Any | None = None
    action: Any | None = None

    # Optional hints (never required)
    observed_transitions: list[str] | None = None

    # Optional metadata
    ts: str | None = None
    source: str | None = None


class TraceArtifact(Artifact):
    """
    A trace artifact can be:
      - legacy single-step: given/when/then
      - multi-step: steps=[TraceStep, ...] (including 1-step logs)

    We keep legacy fields for compatibility, but normalize everything into `steps`.
    """

    # --- NEW preferred representation ---
    steps: list[TraceStep] = Field(default_factory=list)

    # --- Legacy single-step representation (kept for compatibility) ---
    given: str | None = None                  # pre-state (IML literal string)
    when: str | None = None                   # action (IML literal string)
    then: str | None = None                   # post-state (IML literal string)

    # Analysis metadata
    iml_validity: TraceIMLValidity = Field(
        default_factory=lambda: TraceIMLValidity(iml_valid=IMLValidity.UNKNOWN, err=None)
    )
    time: str = ""

    @model_validator(mode="after")
    def _normalize(self) -> "TraceArtifact":
        # If explicit steps exist, accept them (including single-step logs).
        if self.steps:
            return self

        # Otherwise, allow legacy single-step fields and lift into steps.
        if self.given is not None or self.when is not None or self.then is not None:
            # Represent (given, when) as one step. If `then` exists, encode it
            # as a second step with state=then (no action), so post-state is explicit.
            first = TraceStep(state=self.given, action=self.when)
            if self.then is None:
                self.steps = [first]
            else:
                self.steps = [first, TraceStep(state=self.then, action=None)]
            return self

        raise ValueError("Trace must have either `steps` or at least one of `given/when/then`.")

    def is_single_step(self) -> bool:
        return len(self.steps) == 1

    def step_ref(self, step_idx: int) -> str:
        """
        Canonical reference for an individual step of this trace, for use with ArtifactMap.

        NOTE:
          - This is intentionally string-based and does not validate that the step is mapped.
          - We *do* validate that step_idx is in-range, to avoid generating dangling refs.
        """
        if self.art_id is None:
            raise ValueError("Cannot build step ref: art_id is None")
        if step_idx < 0 or step_idx >= len(self.steps):
            raise IndexError(
                f"step_idx out of range (step_idx={step_idx}, num_steps={len(self.steps)})"
            )
        return f"trace:{self.art_id}:step:{step_idx}"


# -----------------------------------------------------------------------------
# Concrete artifacts
# -----------------------------------------------------------------------------

class TestTrace(TraceArtifact):
    """
    Test trace formalization.
    """
    __test__ = False  # pytest: don't treat as a test

    name: str
    filepath: str = ""
    language: str | None = None
    contents: str | None = None


class LogTrace(TraceArtifact):
    """
    Log trace formalization.
    """
    filename: str
    contents: str