"""Step model for ATIF trajectories."""

from datetime import datetime
from typing import Any, Literal

from pydantic import BaseModel, Field, field_validator, model_validator

from harbor.models.trajectories.metrics import Metrics
from harbor.models.trajectories.observation import Observation
from harbor.models.trajectories.tool_call import ToolCall


class Step(BaseModel):
    """A single step in the trajectory."""

    step_id: int = Field(
        default=...,
        ge=1,
        description="Ordinal index of the turn (starting from 1)",
    )
    timestamp: str | None = Field(
        default=None,
        description="ISO 8601 timestamp indicating when this step occurred",
    )
    source: Literal["system", "user", "agent"] = Field(
        default=...,
        description="The originator of this step",
    )
    model_name: str | None = Field(
        default=None,
        description=(
            "The specific LLM model used for this turn. Omission implies the model "
            "defined in the root-level agent config."
        ),
    )
    reasoning_effort: str | float | None = Field(
        default=None,
        description="Qualitative or quantitative measure of effort",
    )
    message: str = Field(
        default=...,
        description="The dialogue message (can be empty string)",
    )
    reasoning_content: str | None = Field(
        default=None,
        description="The agent's explicit internal reasoning",
    )
    tool_calls: list[ToolCall] | None = Field(
        default=None,
        description="Array of structured objects for the agent's actions",
    )
    observation: Observation | None = Field(
        default=None,
        description="Environment feedback/result after actions or system events",
    )
    metrics: Metrics | None = Field(
        default=None,
        description="LLM operational and confidence data for this step",
    )
    extra: dict[str, Any] | None = Field(
        default=None,
        description="Custom step-level metadata",
    )

    model_config = {"extra": "forbid"}

    @field_validator("timestamp")
    @classmethod
    def validate_timestamp(cls, v: str | None) -> str | None:
        """Validate that timestamp is a valid ISO 8601 string."""
        if v is not None:
            try:
                datetime.fromisoformat(v.replace("Z", "+00:00"))
            except ValueError as e:
                raise ValueError(f"Invalid ISO 8601 timestamp: {e}")
        return v

    @model_validator(mode="after")
    def validate_agent_only_fields(self) -> "Step":
        """Validate that certain fields are only present for agent steps."""
        if self.source != "agent":
            agent_only_fields = [
                "model_name",
                "reasoning_effort",
                "reasoning_content",
                "tool_calls",
                "metrics",
            ]
            for field in agent_only_fields:
                if getattr(self, field) is not None:
                    raise ValueError(
                        f"Field '{field}' is only applicable when source is 'agent', "
                        f"but source is '{self.source}'"
                    )
        return self
