"""Type definitions for Harbor ATIF v1.6 trajectory format."""

from __future__ import annotations

from pydantic import BaseModel


class ATIFAgent(BaseModel):
    name: str
    version: str
    model_name: str | None = None


class ATIFToolCall(BaseModel):
    tool_call_id: str
    function_name: str
    arguments: dict = {}


class SubagentTrajectoryRef(BaseModel):
    session_id: str
    trajectory_path: str | None = None


class ATIFObservationResult(BaseModel):
    source_call_id: str | None = None
    content: str | None = None
    subagent_trajectory_ref: list[SubagentTrajectoryRef] | None = None


class ATIFObservation(BaseModel):
    results: list[ATIFObservationResult] = []


class ATIFMetrics(BaseModel):
    prompt_tokens: int | None = None
    completion_tokens: int | None = None
    cached_tokens: int | None = None
    cost_usd: float | None = None


class ATIFFinalMetrics(BaseModel):
    total_prompt_tokens: int | None = None
    total_completion_tokens: int | None = None
    total_cached_tokens: int | None = None
    total_cost_usd: float | None = None
    total_steps: int | None = None


class ATIFStep(BaseModel):
    step_id: int
    timestamp: str | None = None
    source: str  # "user", "agent", "system"
    message: str | None = None
    model_name: str | None = None
    reasoning_content: str | None = None
    tool_calls: list[ATIFToolCall] | None = None
    observation: ATIFObservation | None = None
    metrics: ATIFMetrics | None = None


class ATIFTrajectory(BaseModel):
    schema_version: str = "ATIF-v1.6"
    session_id: str
    agent: ATIFAgent
    steps: list[ATIFStep] = []
    final_metrics: ATIFFinalMetrics | None = None
    notes: str | None = None
