"""Trajectory model for ATIF (Agent Trajectory Interchange Format)."""

from typing import Any, Literal

from pydantic import BaseModel, Field, model_validator

from harbor.models.trajectories.agent import Agent
from harbor.models.trajectories.final_metrics import FinalMetrics
from harbor.models.trajectories.step import Step


class Trajectory(BaseModel):
    """Agent Trajectory in ATIF (Agent Trajectory Interchange Format)."""

    schema_version: Literal["ATIF-v1.0", "ATIF-v1.1", "ATIF-v1.2", "ATIF-v1.3", "ATIF-v1.4"] = Field(
        default="ATIF-v1.4",
        description="String defining ATIF compatibility",
    )
    session_id: str = Field(
        default=...,
        description="Unique identifier for the entire agent run",
    )
    agent: Agent = Field(
        default=...,
        description="Object specifying the agent configuration",
    )
    steps: list[Step] = Field(
        default=...,
        min_length=1,
        description="Array of step objects representing the complete interaction history",
    )
    notes: str | None = Field(
        default=None,
        description="Custom information, design notes, or explanations",
    )
    final_metrics: FinalMetrics | None = Field(
        default=None,
        description="Summary metrics for the entire trajectory",
    )
    extra: dict[str, Any] | None = Field(
        default=None,
        description="Custom root-level metadata",
    )

    model_config = {"extra": "forbid"}

    def to_json_dict(self, exclude_none: bool = True) -> dict[str, Any]:
        """Export trajectory to a dictionary suitable for JSON serialization.

        Args:
            exclude_none: If True, exclude fields with None values from output.

        Returns:
            Dictionary representation of the trajectory.
        """
        return self.model_dump(exclude_none=exclude_none, mode="json")

    @model_validator(mode="after")
    def validate_step_ids(self) -> "Trajectory":
        """Validate that step_ids are sequential starting from 1."""
        for i, step in enumerate(self.steps):
            expected_step_id = i + 1
            if step.step_id != expected_step_id:
                raise ValueError(
                    f"steps[{i}].step_id: expected {expected_step_id} "
                    f"(sequential from 1), got {step.step_id}"
                )
        return self

    @model_validator(mode="after")
    def validate_tool_call_references(self) -> "Trajectory":
        """Validate that observation source_call_ids reference valid tool_call_ids."""
        for step in self.steps:
            if step.observation is None:
                continue

            # Collect all tool_call_ids from this step
            tool_call_ids = set()
            if step.tool_calls:
                tool_call_ids = {tc.tool_call_id for tc in step.tool_calls}

            # Check that source_call_ids reference valid tool_call_ids
            for result in step.observation.results:
                if result.source_call_id is not None:
                    if result.source_call_id not in tool_call_ids:
                        raise ValueError(
                            f"Observation result references source_call_id "
                            f"'{result.source_call_id}' which is not found in "
                            f"step {step.step_id}'s tool_calls"
                        )
        return self
