"""Pydantic types mirroring the Agent Orchestration API schema."""

from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import TYPE_CHECKING

from pydantic import BaseModel, ConfigDict, Field, PrivateAttr

if TYPE_CHECKING:
    from .connectors import Exporter, Importer


# --- Enums ---


class RunStatus(str, Enum):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    TIMEOUT = "timeout"
    CANCELLED = "cancelled"


# --- Common ---


class HealthResponse(BaseModel):
    status: str


class CancelRunResponse(BaseModel):
    status: str


# --- Images ---


class RegistryImage(BaseModel):
    name: str
    tags: list[str]


# --- Agents ---


class AgentCreate(BaseModel):
    image: str = Field(..., description="Agent image reference (e.g. example-agent:latest)")


class Agent(BaseModel):
    id: str
    image: str
    created_at: datetime


# --- Connector specs ---


class RunSnapshotSourceSpec(BaseModel):
    target: str = "run_snapshot"
    include: str | list[str] = Field(default_factory=lambda: ["/workspace/**"])
    exclude: str | list[str] = Field(default_factory=list)
    run_id: str | None = None


class DataBallSourceSpec(BaseModel):
    target: str = "data_ball"
    data: str


class S3DumpDrainSpec(BaseModel):
    target: str = "s3_dump"
    bucket: str | None = None
    prefix: str | None = None


class AgentWorkspaceDrainSpec(BaseModel):
    target: str = "agent_workspace"


class ExportConnectorSpec(BaseModel):
    source: RunSnapshotSourceSpec | None = None
    drain: S3DumpDrainSpec


class ImportConnectorSpec(BaseModel):
    source: RunSnapshotSourceSpec | DataBallSourceSpec | None = None
    drain: AgentWorkspaceDrainSpec


# --- Runs ---


class RunResult(BaseModel):
    run_id: str
    status: str
    output: str | None = None


class WorkflowIds(BaseModel):
    infra: str
    agent: str


class OutputFile(BaseModel):
    path: str
    presigned_url: str | None = None


class FileOutput(BaseModel):
    delivery: str
    files: list[OutputFile]


class RunCreate(BaseModel):
    agent_id: str
    prompt: str
    model: str = Field(default="gpt-5.4", description="OpenAI model to use")
    timeout_seconds: int = Field(default=3600, ge=60, le=86400)
    import_connectors: list[ImportConnectorSpec] = Field(default_factory=list)
    export_connectors: list[ExportConnectorSpec] = Field(default_factory=list)


class RunResumeRequest(BaseModel):
    prompt: str
    model: str = Field(default="gpt-5.4", description="OpenAI model to use")
    timeout_seconds: int = Field(default=3600, ge=60, le=86400)
    import_connectors: list[ImportConnectorSpec] = Field(default_factory=list)
    export_connectors: list[ExportConnectorSpec] = Field(default_factory=list)


class Run(BaseModel):
    model_config = ConfigDict(populate_by_name=True, arbitrary_types_allowed=True)

    id: str = Field(alias="run_id")
    sandbox: str = Field(alias="agent_id")
    parent: str | None = Field(default=None, alias="parent_run_id")
    prompt: str
    model: str
    status: RunStatus
    created_at: datetime
    started_at: datetime | None = None
    completed_at: datetime | None = None
    result: RunResult | None = None
    file_outputs: list[FileOutput] = Field(default_factory=list)
    error: str | None = None
    workflow_ids: WorkflowIds | None = None

    _resource: object = PrivateAttr(default=None)

    def fork(
        self,
        prompt: str,
        *,
        model: str | None = None,
        imports: list[Importer] | None = None,
        exports: list[Exporter] | None = None,
    ) -> Run:
        from .resources.runs import RunsResource
        resource: RunsResource = self._resource  # type: ignore[assignment]
        return resource._fork(
            self.id,
            prompt,
            model=model or self.model,
            imports=imports or [],
            exports=exports or [],
        )

    def refresh(self) -> None:
        """Re-fetch this run from the API, updating all fields in place."""
        from .resources.runs import RunsResource
        resource: RunsResource = self._resource  # type: ignore[assignment]
        updated = resource.get(self.id)
        for field_name in self.model_fields:
            setattr(self, field_name, getattr(updated, field_name))


class RunStatusResponse(BaseModel):
    run_id: str
    status: RunStatus


# --- Checkpoints ---


class ChatHistoryEntry(BaseModel):
    model_config = ConfigDict(extra="allow")


class CheckpointResponse(BaseModel):
    checkpoint_id: str
    index: int
    layer_sha: str
    s3_path: str
    created_at: datetime
    download_url: str | None = None
    chat_history: list[ChatHistoryEntry] | None = None


# --- Validation errors ---


class ValidationErrorItem(BaseModel):
    loc: list[str | int]
    msg: str
    type: str


class HTTPValidationError(BaseModel):
    detail: list[ValidationErrorItem] = Field(default_factory=list)
