"""API response models for the viewer."""

from datetime import datetime
from typing import Any, Generic, TypeVar
from uuid import UUID

from pydantic import BaseModel

T = TypeVar("T")


class PaginatedResponse(BaseModel, Generic[T]):
    """Paginated response wrapper."""

    items: list[T]
    total: int
    page: int
    page_size: int
    total_pages: int


class EvalSummary(BaseModel):
    """Summary of metrics for an agent/model/dataset combination."""

    metrics: list[dict[str, Any]] = []


class JobSummary(BaseModel):
    """Summary of a job for list views."""

    name: str
    id: UUID | None = None
    started_at: datetime | None = None
    finished_at: datetime | None = None
    n_total_trials: int = 0
    n_completed_trials: int = 0
    n_errors: int = 0
    datasets: list[str] = []
    agents: list[str] = []
    providers: list[str] = []
    models: list[str] = []
    environment_type: str | None = None
    evals: dict[str, EvalSummary] = {}


class TaskSummary(BaseModel):
    """Summary of a task group (agent + model + dataset + task) for list views."""

    task_name: str
    source: str | None = None
    agent_name: str | None = None
    model_provider: str | None = None
    model_name: str | None = None
    n_trials: int = 0
    n_completed: int = 0
    n_errors: int = 0
    exception_types: list[str] = []
    avg_reward: float | None = None
    avg_duration_ms: float | None = None


class TrialSummary(BaseModel):
    """Summary of a trial for list views."""

    name: str
    task_name: str
    id: UUID | None = None
    source: str | None = None
    agent_name: str | None = None
    model_provider: str | None = None
    model_name: str | None = None
    reward: float | None = None
    error_type: str | None = None
    started_at: datetime | None = None
    finished_at: datetime | None = None


class FileInfo(BaseModel):
    """Information about a file in a trial directory."""

    path: str  # Relative path from trial dir
    name: str  # File name
    is_dir: bool
    size: int | None = None  # File size in bytes (None for dirs)


class FilterOption(BaseModel):
    """A filter option with a value and count."""

    value: str
    count: int


class JobFilters(BaseModel):
    """Available filter options for jobs list."""

    agents: list[FilterOption]
    providers: list[FilterOption]
    models: list[FilterOption]


class TaskFilters(BaseModel):
    """Available filter options for tasks list within a job."""

    agents: list[FilterOption]
    providers: list[FilterOption]
    models: list[FilterOption]
    tasks: list[FilterOption]


class ComparisonTask(BaseModel):
    """A task identifier for the comparison grid."""

    source: str | None = None
    task_name: str
    key: str


class ComparisonAgentModel(BaseModel):
    """A job+agent+model identifier for the comparison grid."""

    job_name: str
    agent_name: str | None = None
    model_provider: str | None = None
    model_name: str | None = None
    key: str


class ComparisonCell(BaseModel):
    """A cell in the comparison grid."""

    job_name: str
    avg_reward: float | None = None
    avg_duration_ms: float | None = None
    n_trials: int = 0
    n_completed: int = 0


class ComparisonGridData(BaseModel):
    """Data for the job comparison grid view."""

    tasks: list[ComparisonTask]
    agent_models: list[ComparisonAgentModel]
    cells: dict[str, dict[str, ComparisonCell]]  # task.key -> am.key -> cell
