#!/usr/bin/env python

"""████
 ██    ██    Datature
   ██  ██    Powering Breakthrough AI
     ██

@File    :   schemas.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Custom schemas for Vi SDK extensions to OpenAI API.
"""

from typing import Literal

from pydantic import BaseModel, Field


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

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


class ChatCompletionRequest(BaseModel):
    """Request for chat completion endpoint.

    Supports both OpenAI format (messages) and Vi format (source + user_prompt).
    """

    model: str
    messages: list[ChatMessage] | None = None
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    max_tokens: int | None = Field(default=None, ge=1)
    stream: bool = False
    response_format: dict | None = Field(
        default=None,
        description='Structured output format, e.g., {"type": "json_schema", "json_schema": {...}}',
    )

    # Additional generation parameters (OpenAI-compatible)
    top_p: float | None = Field(
        default=None, ge=0.0, le=1.0, description="Top-p sampling parameter"
    )
    top_k: int | None = Field(
        default=None, ge=1, description="Top-k sampling parameter"
    )
    seed: int | None = Field(
        default=None, description="Random seed for reproducible generation"
    )
    repetition_penalty: float | None = Field(
        default=None, ge=0.0, description="Repetition penalty"
    )
    frequency_penalty: float | None = Field(
        default=None, description="Frequency penalty (OpenAI-style)"
    )
    presence_penalty: float | None = Field(
        default=None, description="Presence penalty (OpenAI-style)"
    )
    stop: str | list[str] | None = Field(default=None, description="Stop sequences")
    logit_bias: dict[str, float] | None = Field(
        default=None, description="Logit bias adjustments"
    )

    # Generic generation config (for advanced parameters)
    generation_config: dict | None = Field(
        default=None,
        description="Additional generation parameters as a dictionary. These will be merged with individual parameters.",
    )

    # Vi SDK format support
    source: str | None = Field(
        default=None,
        description=(
            "Image source (Vi SDK format). Can be a file path, URL, or base64-encoded image string"
            "(e.g., 'data:image/jpeg;base64,/9j/...'). Use this OR messages."
        ),
    )
    user_prompt: str | None = Field(
        default=None,
        description="Text prompt (Vi SDK format). Use this OR messages.",
    )


class ChatCompletionStreamChoice(BaseModel):
    """A single choice in a streaming chat completion chunk."""

    index: int
    delta: dict
    finish_reason: Literal["stop", "length", "tool_calls", "content_filter"] | None


class ChatCompletionStreamChunk(BaseModel):
    """A chunk in a streaming chat completion response."""

    id: str
    object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
    created: int
    model: str
    choices: list[ChatCompletionStreamChoice]


class ModelObject(BaseModel):
    """Metadata about a model."""

    id: str
    object: Literal["model"] = "model"
    owned_by: str = "datature"
    created: int | None = None


class ModelListResponse(BaseModel):
    """Response from list models endpoint."""

    object: Literal["list"] = "list"
    data: list[ModelObject]


class ErrorResponse(BaseModel):
    """Error response."""

    error: dict[str, str]
