"""LLMHosts proxy data models -- Pydantic v2 schemas for OpenAI and Anthropic wire formats.

Covers request/response envelopes for both synchronous and streaming modes of
the ``/v1/chat/completions`` (OpenAI) and ``/v1/messages`` (Anthropic) endpoints,
plus a *unified* internal representation used by the dispatcher.
"""

from __future__ import annotations

import time
import uuid
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field

# ============================================================================
# OpenAI dialect
# ============================================================================


class FunctionCall(BaseModel):
    """OpenAI function call inside a tool_call."""

    model_config = ConfigDict(extra="allow")

    name: str
    arguments: str


class ToolCall(BaseModel):
    """OpenAI tool call attached to an assistant message."""

    model_config = ConfigDict(extra="allow")

    id: str
    type: Literal["function"] = "function"
    function: FunctionCall


class ChatMessage(BaseModel):
    """A single message in an OpenAI chat conversation."""

    model_config = ConfigDict(extra="allow")

    role: Literal["system", "user", "assistant", "tool", "function"]
    content: str | list[dict[str, Any]] | None = None
    name: str | None = None
    tool_calls: list[ToolCall] | None = None
    tool_call_id: str | None = None


class ResponseFormat(BaseModel):
    """OpenAI response_format parameter."""

    model_config = ConfigDict(extra="allow")

    type: str = "text"


class ChatCompletionRequest(BaseModel):
    """OpenAI ``POST /v1/chat/completions`` request body."""

    model_config = ConfigDict(extra="allow")

    model: str
    messages: list[ChatMessage]
    stream: bool = False
    temperature: float | None = None
    max_tokens: int | None = None
    top_p: float | None = None
    stop: str | list[str] | None = None
    tools: list[dict[str, Any]] | None = None
    tool_choice: str | dict[str, Any] | None = None
    response_format: ResponseFormat | None = None
    n: int | None = 1
    frequency_penalty: float | None = None
    presence_penalty: float | None = None
    user: str | None = None


class Usage(BaseModel):
    """Token usage information for OpenAI responses."""

    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0


class ChatCompletionChoice(BaseModel):
    """A single choice in a non-streaming OpenAI response."""

    index: int = 0
    message: ChatMessage
    finish_reason: str | None = "stop"


class ChatCompletionResponse(BaseModel):
    """OpenAI ``POST /v1/chat/completions`` non-streaming response."""

    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:24]}")
    object: str = "chat.completion"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str = ""
    choices: list[ChatCompletionChoice] = Field(default_factory=list)
    usage: Usage = Field(default_factory=Usage)


class DeltaMessage(BaseModel):
    """Partial message used in streaming chunks."""

    role: str | None = None
    content: str | None = None
    tool_calls: list[ToolCall] | None = None


class ChatCompletionChunkChoice(BaseModel):
    """A single choice inside a streaming chunk."""

    index: int = 0
    delta: DeltaMessage = Field(default_factory=DeltaMessage)
    finish_reason: str | None = None


class ChatCompletionChunk(BaseModel):
    """OpenAI streaming chunk (one SSE ``data:`` frame)."""

    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:24]}")
    object: str = "chat.completion.chunk"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str = ""
    choices: list[ChatCompletionChunkChoice] = Field(default_factory=list)


# ============================================================================
# Anthropic dialect
# ============================================================================


class AnthropicContentBlock(BaseModel):
    """Anthropic content block (text, image, tool_use, etc.)."""

    model_config = ConfigDict(extra="allow")

    type: str = "text"
    text: str | None = None
    source: dict[str, Any] | None = None
    id: str | None = None
    name: str | None = None
    input: dict[str, Any] | None = None


class AnthropicMessage(BaseModel):
    """A single message in an Anthropic conversation."""

    model_config = ConfigDict(extra="allow")

    role: Literal["user", "assistant"]
    content: str | list[AnthropicContentBlock]


class AnthropicMetadata(BaseModel):
    """Anthropic request metadata."""

    model_config = ConfigDict(extra="allow")

    user_id: str | None = None


class AnthropicRequest(BaseModel):
    """Anthropic ``POST /v1/messages`` request body."""

    model_config = ConfigDict(extra="allow")

    model: str
    messages: list[AnthropicMessage]
    max_tokens: int = 4096
    system: str | None = None
    stream: bool = False
    temperature: float | None = None
    top_p: float | None = None
    top_k: int | None = None
    stop_sequences: list[str] | None = None
    metadata: AnthropicMetadata | None = None


class AnthropicUsage(BaseModel):
    """Token counts in Anthropic responses."""

    input_tokens: int = 0
    output_tokens: int = 0


class AnthropicResponse(BaseModel):
    """Anthropic ``POST /v1/messages`` non-streaming response."""

    id: str = Field(default_factory=lambda: f"msg_{uuid.uuid4().hex[:24]}")
    type: str = "message"
    role: str = "assistant"
    content: list[AnthropicContentBlock] = Field(default_factory=list)
    model: str = ""
    stop_reason: str | None = "end_turn"
    stop_sequence: str | None = None
    usage: AnthropicUsage = Field(default_factory=AnthropicUsage)


class AnthropicStreamEvent(BaseModel):
    """A single Anthropic SSE event.

    Anthropic streaming uses typed events::

        event: message_start
        data: {"type": "message_start", "message": {...}}

        event: content_block_delta
        data: {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hi"}}

        event: message_stop
        data: {"type": "message_stop"}
    """

    model_config = ConfigDict(extra="allow")

    type: Literal[
        "message_start",
        "content_block_start",
        "content_block_delta",
        "content_block_stop",
        "message_delta",
        "message_stop",
        "ping",
    ]
    message: dict[str, Any] | None = None
    index: int | None = None
    content_block: dict[str, Any] | None = None
    delta: dict[str, Any] | None = None
    usage: dict[str, Any] | None = None


# ============================================================================
# Unified internal representation
# ============================================================================


class UnifiedMessage(BaseModel):
    """Internal message representation used between dialect adapters and the dispatcher."""

    role: str
    content: str | list[dict[str, Any]] | None = None
    name: str | None = None
    tool_calls: list[dict[str, Any]] | None = None
    tool_call_id: str | None = None


class UnifiedRequest(BaseModel):
    """Dialect-agnostic request representation.

    Both the OpenAI and Anthropic route handlers convert their incoming
    request to this format before handing it to :class:`BackendDispatcher`.
    """

    model: str
    messages: list[UnifiedMessage]
    stream: bool = False
    temperature: float | None = None
    max_tokens: int | None = None
    top_p: float | None = None
    stop: list[str] | None = None
    tools: list[dict[str, Any]] | None = None
    extra: dict[str, Any] = Field(default_factory=dict)


class UnifiedResponse(BaseModel):
    """Dialect-agnostic response representation.

    The dispatcher returns this; route handlers translate it back to the
    caller's dialect.
    """

    content: str = ""
    finish_reason: str = "stop"
    model: str = ""
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    tool_calls: list[dict[str, Any]] | None = None
    extra: dict[str, Any] | None = None  # Dispatcher metadata (e.g. fallback_info)


# ============================================================================
# /v1/models list
# ============================================================================


class ModelInfo(BaseModel):
    """Single model entry for ``GET /v1/models``."""

    id: str
    object: str = "model"
    created: int = Field(default_factory=lambda: int(time.time()))
    owned_by: str = "llmhosts"


class ModelListResponse(BaseModel):
    """``GET /v1/models`` response."""

    object: str = "list"
    data: list[ModelInfo] = Field(default_factory=list)
