from __future__ import annotations

from datetime import datetime
from decimal import Decimal
from typing import ClassVar
from uuid import UUID

from pydantic import BaseModel, ConfigDict, Field
from schemas.agent.runtime_models import RouterAgentOutput, WorkerAgentOutputRich

from ..agent import AgentType, ToolAgentOutput


class UserMessageAttachments(BaseModel):
    model_config: ClassVar[ConfigDict] = ConfigDict(extra="allow")

    bucket: str
    path: str
    mime_type: str


class AgentChatMessageMetadata(BaseModel):
    model_config: ClassVar[ConfigDict] = ConfigDict(extra="allow")
    run_id: str
    agent_type: AgentType | None = None
    user_message_attachments: UserMessageAttachments | None = None
    router_agent_output: RouterAgentOutput | None = None
    tool_agent_output: ToolAgentOutput | None = None
    worker_agent_output: WorkerAgentOutputRich | None = None


class AgentChatMessage(BaseModel):
    """Canonical schema aligned with `messages` table columns."""

    model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid")

    id: UUID
    seq: int
    role: str
    content: str
    model_code: str | None = None
    tool_name: str | None = None
    input_tokens: int = Field(default=0, ge=0)
    output_tokens: int = Field(default=0, ge=0)
    cost: Decimal = Field(default=Decimal("0"))
    latency_ms: int | None = Field(default=None, ge=0)
    metadata: AgentChatMessageMetadata | dict[str, object] | None = None
    timestamp: datetime