"""Core Pydantic schemas for agent API."""
from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, model_validator

from logos_adk.types import ContentPart


class ContentPartPayload(BaseModel):
    """HTTP-serializable multimodal content block."""

    type: str
    text: str | None = None
    uri: str | None = None
    data: str | None = None
    mime_type: str | None = None
    name: str | None = None
    metadata: dict[str, Any] = {}

    def to_runtime_part(self) -> ContentPart:
        """Convert to runtime ContentPart."""
        return ContentPart(
            type=self.type,
            text=self.text,
            uri=self.uri,
            data=self.data,
            mime_type=self.mime_type,
            name=self.name,
            metadata=dict(self.metadata),
        )

    @classmethod
    def from_runtime_part(cls, part: ContentPart) -> "ContentPartPayload":
        """Create from runtime ContentPart."""
        return cls(
            type=part.type,
            text=part.text,
            uri=part.uri,
            data=part.data,
            mime_type=part.mime_type,
            name=part.name,
            metadata=dict(part.metadata),
        )


class TaskRunSection(BaseModel):
    """A schema-driven result section snapshot persisted with a task run."""

    id: str
    kind: Literal["summary", "bullets", "next_steps", "full_text"]
    title: str
    maxItems: int | None = None


class TaskRunLayout(BaseModel):
    """Snapshot of the result layout used for a task run."""

    title: str | None = None
    sections: list[TaskRunSection] = []


class TaskRunPayload(BaseModel):
    """Client-provided task run snapshot for server-side session continuity."""

    id: str
    sessionId: str | None = None
    agentName: str
    profileId: str
    profileTitle: str
    profileBadge: str
    scenarioId: str | None = None
    scenarioTitle: str | None = None
    prompt: str
    fieldValues: dict[str, str] = {}
    launchTitle: str | None = None
    resultLayout: TaskRunLayout | None = None
    status: Literal["running", "completed", "failed", "paused"] = "running"
    startedAt: str
    completedAt: str | None = None
    runId: str | None = None
    requestId: str | None = None
    traceId: str | None = None


class TaskRunRecord(BaseModel):
    """Persisted metadata for a single user task run."""

    id: str
    session_id: str
    agent_name: str
    profile_id: str
    profile_title: str
    profile_badge: str
    scenario_id: str | None = None
    scenario_title: str | None = None
    prompt: str
    field_values: dict[str, str] = {}
    launch_title: str | None = None
    result_layout: TaskRunLayout | None = None
    status: Literal["running", "completed", "failed", "paused"]
    started_at: str
    completed_at: str | None = None
    run_id: str | None = None
    request_id: str | None = None
    trace_id: str | None = None


class AgentRequest(BaseModel):
    """Request body for /run and /stream endpoints."""

    prompt: str | None = None
    prompt_parts: list[ContentPartPayload] | None = None
    session_id: str | None = None
    task_run: TaskRunPayload | None = None

    @model_validator(mode="after")
    def validate_prompt(self) -> "AgentRequest":
        """Validate that prompt or prompt_parts is provided."""
        if not self.prompt and not self.prompt_parts:
            raise ValueError("prompt or prompt_parts is required")
        return self


class BatchAgentRequest(BaseModel):
    """Request body for /batch endpoint."""

    prompts: list[str]
    session_ids: list[str | None] | None = None
    max_concurrency: int = 4


class UIStreamRequest(BaseModel):
    """Request body for /stream-ui endpoint."""

    prompt: str | None = None
    prompt_parts: list[ContentPartPayload] | None = None
    session_id: str | None = None
    correlation_id: str | None = None
    task_run: TaskRunPayload | None = None
    task_run_id: str | None = None
    user_input: str | None = None
    continue_from_node: str | None = None
    continue_correlation_id: str | None = None

    @model_validator(mode="after")
    def validate_prompt(self) -> "UIStreamRequest":
        """Validate that prompt or prompt_parts is provided."""
        if not self.prompt and not self.prompt_parts:
            raise ValueError("prompt or prompt_parts is required")
        return self


# Alias for backward compatibility
StreamUIRequest = UIStreamRequest


class ResumeRequest(BaseModel):
    """Request body for /resume endpoint to continue paused execution."""

    prompt: str | None = None
    prompt_parts: list[ContentPartPayload] | None = None
    session_id: str | None = None
    correlation_id: str | None = None
    task_run: TaskRunPayload | None = None
    task_run_id: str | None = None
    node_name: str | None = None
    user_input: str | None = None
    continue_from_node: str | None = None
    continue_correlation_id: str | None = None

    @model_validator(mode="after")
    def validate_prompt(self) -> "ResumeRequest":
        """Validate that prompt or prompt_parts is provided."""
        if not self.prompt and not self.prompt_parts:
            raise ValueError("prompt or prompt_parts is required")
        return self


class ToolCallResponse(BaseModel):
    """Serialized tool call in response."""

    id: str
    name: str
    arguments: dict


class UsageResponse(BaseModel):
    """Token usage in response."""

    input_tokens: int
    output_tokens: int


class AgentResponse(BaseModel):
    """Response body for /run endpoint."""

    content: str
    content_parts: list[ContentPartPayload] = []
    tool_calls: list[ToolCallResponse]
    usage: UsageResponse
    run_id: str | None = None
    trace_id: str | None = None
    request_id: str | None = None
    session_id: str | None = None
    workspace_id: str | None = None
    prompt_name: str | None = None
    prompt_version: str | None = None
    prompt_variant: str | None = None
    experiment_id: str | None = None
    model: str | None = None
    config_fingerprint: str | None = None


class BatchAgentResponse(BaseModel):
    """Response body for /batch endpoint."""

    items: list[AgentResponse]


__all__ = [
    "ContentPartPayload",
    "TaskRunSection",
    "TaskRunLayout",
    "TaskRunPayload",
    "TaskRunRecord",
    "AgentRequest",
    "BatchAgentRequest",
    "UIStreamRequest",
    "StreamUIRequest",
    "ResumeRequest",
    "ToolCallResponse",
    "UsageResponse",
    "AgentResponse",
    "BatchAgentResponse",
]
