"""Schemas for uploading agent trial results."""

from __future__ import annotations

from typing import Any, Literal

from openai.types.chat import ChatCompletionMessageParam
from pydantic import BaseModel, Field


class CriterionRatingIn(BaseModel):
    """A single LLM-judge rating for a named criterion."""

    label: str = Field(description="Criterion label (max 25 chars)")
    description: str = Field(default="", description="Criterion description")
    rating: Literal["YES", "NO"] = Field(description="YES or NO judgment")
    reasoning: str = Field(default="", description="LLM judge explanation")


class AgentTrialResult(BaseModel):
    """A single agent trial result for upload."""

    task_name: str = Field(description="Maps to collection variable via task_name_column")
    messages: list[ChatCompletionMessageParam] = Field(
        default_factory=list, description="Agent trace in OpenAI message format"
    )
    reward: float | None = Field(default=None, description="Primary reward score (0.0-1.0)")
    steps: int | None = Field(default=None, description="Number of agent steps")
    cost_usd: float | None = Field(default=None, description="Total cost in USD")
    duration_seconds: float | None = Field(default=None, description="Trial duration in seconds")
    input_tokens: int | None = Field(default=None, description="Total input tokens")
    output_tokens: int | None = Field(default=None, description="Total output tokens")
    cached_tokens: int | None = Field(default=None, description="Total cached input tokens")
    error: str | None = Field(default=None, description="Error message if trial failed")
    metadata: dict[str, Any] | None = Field(default=None, description="Additional metadata")
    trajectory: dict[str, Any] | None = Field(
        default=None, description="Raw ATIF trajectory dict (validated by backend)"
    )
    criterion_ratings: list[CriterionRatingIn] = Field(default_factory=list, description="Per-criterion ratings")


class AgentResultsUploadRequest(BaseModel):
    """Request payload for uploading agent trial results."""

    results: list[AgentTrialResult]
    epoch: int = Field(default=1, ge=1)
    task_name_column: str
    evaluate: bool = Field(
        default=True,
        description="Whether to auto-trigger trajectory evaluation after upload.",
    )


class AgentResultsUploadResponse(BaseModel):
    """Response from agent results upload."""

    created_responses: int
    created_ratings: int
    errors: list[str] = Field(default_factory=list)
    pending_evaluations: int = 0
