from datetime import datetime
from typing import Any
from uuid import UUID

from ninja import Schema


class AgentStepIn(Schema):
    code: str
    skill_version_id: int
    input_mapping: dict[str, str]


class AgentCreateIn(Schema):
    code: str
    name: str
    description: str = ""
    model_deployment_id: int
    system_prompt: str = ""
    input_schema: dict[str, Any] = {"type": "object"}
    example_input: dict[str, Any] = {}
    temperature_override: float | None = None
    max_output_tokens_override: int | None = None
    steps: list[AgentStepIn]


class AgentDefinitionUpdateIn(Schema):
    name: str | None = None
    description: str | None = None


class AgentDraftUpdateIn(Schema):
    model_deployment_id: int | None = None
    system_prompt: str | None = None
    input_schema: dict[str, Any] | None = None
    example_input: dict[str, Any] | None = None
    temperature_override: float | None = None
    max_output_tokens_override: int | None = None
    steps: list[AgentStepIn] | None = None


class AgentStepOut(Schema):
    id: int
    code: str
    order: int
    skill_version_id: int
    skill_name: str
    skill_version: int
    skill_status: str
    input_mapping: dict[str, str]


class AgentVersionOut(Schema):
    id: int
    version: int
    status: str
    model_deployment_id: int
    model_display_name: str
    system_prompt: str
    input_schema: dict[str, Any]
    example_input: dict[str, Any]
    temperature_override: float | None
    max_output_tokens_override: int | None
    steps: list[AgentStepOut]
    created_at: datetime
    updated_at: datetime
    published_at: datetime | None


class AgentDefinitionOut(Schema):
    id: int
    code: str
    name: str
    description: str
    active_version: int | None
    draft_version: int | None
    updated_at: datetime


class AgentDetailOut(AgentDefinitionOut):
    versions: list[AgentVersionOut]


class AgentRunIn(Schema):
    version: str | int = "draft"
    input: dict[str, Any]


class AgentRunAttemptOut(Schema):
    id: int
    attempt_type: str
    status: str
    raw_output: str
    validation_error: str
    input_tokens: int
    output_tokens: int
    duration_ms: int
    created_at: datetime


class AgentRunStepOut(Schema):
    id: int
    code: str
    name: str
    order: int
    status: str
    input_payload: dict[str, Any]
    output_payload: Any | None
    input_tokens: int
    output_tokens: int
    duration_ms: int
    error_code: str
    error_message: str
    attempts: list[AgentRunAttemptOut]


class AgentRunSummaryOut(Schema):
    public_id: UUID
    agent_id: int
    agent_name: str
    agent_version: int
    status: str
    sandbox: bool
    current_step: str
    total_steps: int
    input_tokens: int
    output_tokens: int
    duration_ms: int
    error_code: str
    error_message: str
    created_at: datetime
    started_at: datetime | None
    finished_at: datetime | None


class AgentRunDetailOut(AgentRunSummaryOut):
    input_payload: dict[str, Any]
    final_output: Any | None
    config_snapshot: dict[str, Any]
    cancellation_requested_at: datetime | None
    steps: list[AgentRunStepOut]