"""State models for the agent runtime."""
from __future__ import annotations

from dataclasses import asdict, dataclass, field
from enum import Enum
from typing import Any, Optional


class StepStatus(str, Enum):
    PENDING = "pending"
    RUNNING = "running"
    DONE = "done"
    FAILED = "failed"


class RunStatus(str, Enum):
    RUNNING = "running"
    DONE = "done"
    FAILED = "failed"


@dataclass
class PlanStep:
    """A single step within an agent plan."""

    step_id: str
    kind: str  # "shell" | "write_file" | "read_file" | "finish" | "flaky"
    description: str
    payload: dict[str, Any]
    status: StepStatus = StepStatus.PENDING
    attempts: int = 0  # how many times this step has been attempted (0 = not yet started)

    def to_dict(self) -> dict[str, Any]:
        d = asdict(self)
        d["status"] = self.status.value
        return d


@dataclass
class StepResult:
    """Result of executing a single plan step."""

    step_id: str
    status: StepStatus
    output_text: Optional[str] = None
    error: Optional[str] = None
    metadata: dict[str, Any] = field(default_factory=dict)
    started_at: str = ""
    finished_at: str = ""
    attempt: int = 1               # which attempt produced this result (1-based)
    failure_class: Optional[str] = None  # set on FAILED results

    def to_dict(self) -> dict[str, Any]:
        return {
            "step_id": self.step_id,
            "status": self.status.value,
            "output_text": self.output_text,
            "error": self.error,
            "metadata": self.metadata,
            "started_at": self.started_at,
            "finished_at": self.finished_at,
            "attempt": self.attempt,
            "failure_class": self.failure_class,
        }


@dataclass
class Plan:
    """Ordered list of steps to accomplish a goal."""

    steps: list[PlanStep] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {"steps": [s.to_dict() for s in self.steps]}


@dataclass
class AgentState:
    """Full runtime state of a single agent run."""

    run_id: str
    goal: str
    plan: Plan
    current_step_id: Optional[str]
    history: list[StepResult]
    started_at: str
    updated_at: str
    status: RunStatus
    trace_id: Optional[str] = None

    # --- History access helpers ---

    def get_last_output(self) -> Optional[str]:
        """Return output_text of the most recently completed step, or None."""
        for result in reversed(self.history):
            if result.output_text is not None:
                return result.output_text
        return None

    def get_step_output(self, step_id: str) -> Optional[str]:
        """Return output_text for a specific step_id, or None if not found."""
        for result in self.history:
            if result.step_id == step_id:
                return result.output_text
        return None

    def list_outputs(self) -> list[tuple[str, Optional[str]]]:
        """Return [(step_id, output_text)] for all history entries."""
        return [(r.step_id, r.output_text) for r in self.history]

    def to_dict(self) -> dict[str, Any]:
        return {
            "run_id": self.run_id,
            "trace_id": self.trace_id,
            "goal": self.goal,
            "plan": self.plan.to_dict(),
            "current_step_id": self.current_step_id,
            "history": [r.to_dict() for r in self.history],
            "started_at": self.started_at,
            "updated_at": self.updated_at,
            "status": self.status.value,
        }