"""Pydantic V2 API request/response schemas."""

from pydantic import BaseModel, ConfigDict, Field, field_validator

from llm_inference.types import (
    BackendType,
    ChatMessage,
    Choice,
    ModelInfo,
    StreamChoice,
    Usage,
)

# =============================================================================
# Completion Schemas
# =============================================================================


class CompletionRequest(BaseModel):
    """Request schema for chat completions."""

    model_config = ConfigDict(extra="forbid")

    messages: list[ChatMessage] = Field(
        min_length=1,
        max_length=1000,
        description="List of messages in the conversation (1-1000 messages)",
    )
    model: str = Field(
        description="Model to use for completion",
    )
    temperature: float = Field(
        default=0.7,
        ge=0.0,
        le=2.0,
        description="Sampling temperature",
    )
    max_tokens: int | None = Field(
        default=None,
        ge=1,
        le=1000000,
        description="Maximum tokens to generate (1-1000000)",
    )
    stream: bool = Field(
        default=False,
        description="Enable streaming response",
    )
    backend: BackendType | None = Field(
        default=None,
        description="Backend to use (overrides default)",
    )
    top_p: float | None = Field(
        default=None,
        ge=0.0,
        le=1.0,
        description="Top-p sampling parameter",
    )
    frequency_penalty: float | None = Field(
        default=None,
        ge=-2.0,
        le=2.0,
        description="Frequency penalty",
    )
    presence_penalty: float | None = Field(
        default=None,
        ge=-2.0,
        le=2.0,
        description="Presence penalty",
    )
    stop: list[str] | str | None = Field(
        default=None,
        description="Stop sequences",
    )

    @field_validator("messages")
    @classmethod
    def validate_messages_content(cls, v: list[ChatMessage]) -> list[ChatMessage]:
        """Validate that messages have content (except assistant messages)."""
        for i, msg in enumerate(v):
            # Assistant messages can have empty content (for function calls etc)
            if msg.role != "assistant" and msg.content is None:
                raise ValueError(
                    f"Message at index {i} with role '{msg.role}' cannot have empty content"
                )
        return v


class CompletionResponse(BaseModel):
    """Response schema for chat completions."""

    model_config = ConfigDict(extra="forbid")

    id: str = Field(description="Unique completion ID")
    object: str = Field(default="chat.completion")
    created: int = Field(description="Unix timestamp of creation")
    model: str = Field(description="Model used for completion")
    choices: list[Choice] = Field(description="Completion choices")
    usage: Usage | None = Field(default=None, description="Token usage")
    backend: str = Field(description="Backend that served the request")


class CompletionChunk(BaseModel):
    """Streaming chunk for chat completions."""

    model_config = ConfigDict(extra="forbid")

    id: str = Field(description="Unique completion ID")
    object: str = Field(default="chat.completion.chunk")
    created: int = Field(description="Unix timestamp of creation")
    model: str = Field(description="Model used for completion")
    choices: list[StreamChoice] = Field(description="Streaming choices")


# =============================================================================
# Text Completion Schemas (legacy OpenAI /v1/completions)
# =============================================================================


class TextChoice(BaseModel):
    """One text-completion choice."""

    model_config = ConfigDict(extra="forbid")

    index: int = Field(description="Choice index")
    text: str = Field(description="Generated text")
    finish_reason: str | None = Field(default=None, description="Finish reason")


class TextCompletionRequest(BaseModel):
    """Request schema for text completions (/v1/completions)."""

    model_config = ConfigDict(extra="forbid")

    prompt: str | list[str] = Field(description="Prompt(s) to complete")
    model: str | None = Field(default=None, description="Model (alias) to use")
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    max_tokens: int | None = Field(default=None, ge=1, le=1000000)
    backend: BackendType | None = Field(
        default=None, description="Backend to use (overrides default)"
    )
    top_p: float | None = Field(default=None, ge=0.0, le=1.0)
    frequency_penalty: float | None = Field(default=None, ge=-2.0, le=2.0)
    presence_penalty: float | None = Field(default=None, ge=-2.0, le=2.0)
    stop: list[str] | str | None = Field(default=None, description="Stop sequences")


class TextCompletionResponse(BaseModel):
    """Response schema for text completions."""

    model_config = ConfigDict(extra="forbid")

    id: str = Field(description="Unique completion ID")
    object: str = Field(default="text_completion")
    created: int = Field(description="Unix timestamp of creation")
    model: str = Field(description="Model used for completion")
    choices: list[TextChoice] = Field(description="Completion choices")
    usage: Usage | None = Field(default=None, description="Token usage")
    backend: str = Field(description="Backend that served the request")


# =============================================================================
# Model Management Schemas
# =============================================================================


class ModelListResponse(BaseModel):
    """Response schema for listing models."""

    object: str = Field(default="list")
    data: list[ModelInfo] = Field(description="List of available models")


class ModelLoadRequest(BaseModel):
    """Request schema for loading a model."""

    model: str = Field(description="Model identifier to load")
    backend: BackendType = Field(description="Backend to load the model on")


class ModelLoadResponse(BaseModel):
    """Response schema for model load/unload operations."""

    success: bool = Field(description="Whether the operation succeeded")
    model: str = Field(description="Model identifier")
    backend: str = Field(description="Backend the operation was performed on")
    message: str | None = Field(default=None, description="Optional message")


# =============================================================================
# Health Schemas
# =============================================================================


class BackendHealth(BaseModel):
    """Health status for a single backend."""

    name: str = Field(description="Backend name")
    healthy: bool = Field(description="Whether the backend is healthy")
    message: str | None = Field(default=None, description="Optional status message")


class HealthResponse(BaseModel):
    """Response schema for health check."""

    status: str = Field(description="Overall status: healthy, degraded, unhealthy")
    backends: list[BackendHealth] = Field(description="Per-backend health status")


class ReadinessResponse(BaseModel):
    """Response schema for readiness probe."""

    ready: bool = Field(default=True)


# =============================================================================
# Backend Schemas
# =============================================================================


class BackendListResponse(BaseModel):
    """Response schema for listing backends."""

    backends: list[str] = Field(description="List of available backend names")