"""Chat session Pydantic schemas."""
from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel

from .agent import ContentPartPayload, TaskRunRecord, UsageResponse


class ChatMessage(BaseModel):
    """A single chat message."""

    id: str
    role: Literal["user", "assistant", "system"]
    content: str
    content_parts: list[ContentPartPayload] = []
    timestamp: str
    agent_name: str | None = None
    usage: UsageResponse | None = None
    task_run_id: str | None = None
    run_id: str | None = None
    request_id: str | None = None
    trace_id: str | None = None


class ChatSession(BaseModel):
    """A chat session containing messages."""

    id: str
    name: str
    agent_name: str | None = None
    messages: list[ChatMessage] = []
    task_runs: list[TaskRunRecord] = []
    tokens_by_agent: dict[str, Any] = {}
    created_at: str
    updated_at: str | None = None


class CreateChatSessionRequest(BaseModel):
    """Request to create a new chat session."""

    name: str | None = None
    agent_name: str | None = None


class UpdateChatSessionRequest(BaseModel):
    """Request to update a chat session."""

    name: str | None = None
    messages: list[ChatMessage] | None = None
    task_runs: list[TaskRunRecord] | None = None
    tokens_by_agent: dict[str, Any] | None = None


class ChatSessionsResponse(BaseModel):
    """Response containing list of chat sessions."""

    sessions: list[ChatSession]
    total: int


class TaskRunActionResponse(BaseModel):
    """Response body for task-run actions like cancel."""

    status: Literal["cancel_requested", "canceled", "already_terminal"]
    session_id: str
    task_run_id: str
    task_run_status: Literal["running", "completed", "failed", "paused"]


__all__ = [
    "ChatMessage",
    "ChatSession",
    "CreateChatSessionRequest",
    "UpdateChatSessionRequest",
    "ChatSessionsResponse",
    "TaskRunActionResponse",
]
