"""Data models for trajectory viewer."""

from __future__ import annotations

from pathlib import Path
from typing import Literal

from pydantic import BaseModel

from harbor.models.trajectories.metrics import Metrics
from harbor.models.trajectories.observation import Observation
from harbor.models.trajectories.observation_result import ObservationResult
from harbor.models.verifier.result import VerifierResult

Status = Literal["completed", "running", "unknown"]


# =============================================================================
# Summary Models (for list endpoints) - viewer-specific
# =============================================================================


class JobSummary(BaseModel):
    """Summary info for a job directory."""

    job_id: str
    job_path: Path
    trial_count: int
    completed_count: int  # trials with result.json
    success_count: int  # trials with reward > 0
    failed_count: int  # completed trials with reward == 0 or None
    status: Status
    mtime: float | None = None


class TrialSummary(BaseModel):
    """Summary info for a trial directory."""

    trial_id: str
    trial_path: Path
    task_name: str | None = None
    has_trajectory: bool
    trajectory_mtime: float | None = None
    status: Status
    rewards: dict[str, float] | None = None
    exception_type: str | None = None


# =============================================================================
# Serialization Helpers - convert datetime to ISO string for JSON
# =============================================================================


class TimingResponse(BaseModel):
    """Timing info serialized for API response (datetime → str)."""

    started_at: str | None = None
    finished_at: str | None = None


class ExceptionInfoResponse(BaseModel):
    """Exception information serialized for API response (datetime → str)."""

    exception_type: str
    exception_message: str
    exception_traceback: str
    occurred_at: str  # ISO format string


# =============================================================================
# Viewer-Specific Models
# =============================================================================


class LogsResponse(BaseModel):
    """Log files content for viewer display."""

    trial: str | None = None
    verifier_stdout: str | None = None
    verifier_stderr: str | None = None


class TimingBreakdownResponse(BaseModel):
    """Timing breakdown for all trial phases."""

    environment_setup: TimingResponse | None = None
    agent_setup: TimingResponse | None = None
    agent_execution: TimingResponse | None = None
    verifier: TimingResponse | None = None
    total: TimingResponse | None = None


class TrajectoryMetricsResponse(BaseModel):
    """Aggregated metrics for viewer display.

    Note: This differs from FinalMetrics as it uses different field names
    and is computed from steps when final_metrics is not available.
    """

    cost_usd: float = 0.0
    prompt_tokens: int = 0
    completion_tokens: int = 0
    cached_tokens: int = 0


class ToolCallWithObservation(BaseModel):
    """Tool call paired with its observation results for viewer display.

    This extends the base ToolCall by pairing it with observation results
    for unified rendering in the viewer UI.
    """

    tool_call_id: str
    function_name: str
    arguments: str  # JSON string for display
    observation_results: list[ObservationResult] = []


class StepResponse(BaseModel):
    """Step data formatted for API response.

    This model reformats Step data for viewer display, including:
    - Pairing tool calls with their observation results
    - Separating unmatched observations
    """

    step_id: int
    source: str  # "system" | "user" | "agent"
    timestamp: str | None = None
    message: str
    reasoning_content: str | None = None
    model_name: str | None = None
    reasoning_effort: str | float | None = None
    tool_calls: list[ToolCallWithObservation] | None = None
    observation: Observation | None = None  # Unmatched observations only
    metrics: Metrics | None = None


class TrajectoryResponse(BaseModel):
    """Complete trajectory data for API response.

    This is the main response model returned by:
    - _format_trajectory_info()
    - _format_metadata_only()
    - get_trajectory_data()
    """

    task_name: str | None = None
    session_id: str | None = None
    agent_name: str
    agent_version: str | None = None
    model_name: str | None = None
    schema_version: str | None = None
    steps: list[StepResponse]
    metrics: TrajectoryMetricsResponse
    verifier: VerifierResult | None = None
    notes: str | None = None
    config: str | None = None  # JSON string
    logs: LogsResponse
    exception: str | None = None
    exception_info: ExceptionInfoResponse | None = None
    timing: TimingBreakdownResponse | None = None
    has_trajectory: bool
    parent_session_id: str | None = None
    subagent_trajectories: dict[str, TrajectoryResponse] | None = None
