from __future__ import annotations

import secrets
import time

from pydantic import BaseModel, Field

from orchard.server.models.chat.logprobs import ChatCompletionLogProbs
from orchard.server.models.chat.request import ChatMessage

# --- Constants ---
CHAT_COMPLETION_ID_PREFIX = "chatcmpl-"
CHAT_COMPLETION_OBJECT = "chat.completion"
CHAT_COMPLETION_CHUNK_OBJECT = "chat.completion.chunk"


def generate_chat_completion_id(prefix: str = CHAT_COMPLETION_ID_PREFIX) -> str:
    """Generates a unique identifier string for a completion response."""
    random_part = secrets.token_urlsafe(22)
    return f"{prefix}{random_part}"


def get_current_timestamp() -> int:
    """Returns the current time as a Unix epoch timestamp (seconds)."""
    return int(time.time())


class ChatCompletionChoice(BaseModel):
    """Represents a single generated chat completion choice."""

    index: int = Field(description="The index of this choice.")
    message: ChatMessage = Field(description="The message generated by the model.")
    finish_reason: str | None = Field(
        description="Reason generation stopped (e.g., 'stop', 'length', 'tool_calls')."
    )
    logprobs: ChatCompletionLogProbs.LogProbsContent | None = Field(
        default=None,
        description="Log probability information for the choice.",
    )


class ChatCompletionChunkChoice(BaseModel):
    """A choice from a chunk of a chat completion."""

    delta: ChatMessage = Field(
        description="A chat completion delta generated by streamed model responses."
    )
    finish_reason: str | None = Field(
        description="Reason generation stopped (e.g., 'stop', 'length', 'tool_calls')."
    )
    index: int = Field(description="The index of this choice.")
    logprobs: ChatCompletionLogProbs.LogProbsContent | None = Field(
        default=None,
        description="Log probability information for the choice.",
    )


class ChatCompletionUsage(BaseModel):
    """Provides token usage statistics for the chat completion request."""

    input_tokens: int = Field(
        description="The number of tokens constituting the input prompt(s)."
    )
    output_tokens: int = Field(
        description="The total number of tokens generated across all completion choices."
    )
    total_tokens: int = Field(
        description="The sum of `input_tokens` and `output_tokens`."
    )


# --- Main Response Model ---
class ChatCompletionResponse(BaseModel):
    """Defines the response schema for the chat completion endpoint."""

    id: str = Field(
        default_factory=generate_chat_completion_id,
        description="A unique identifier for the chat completion.",
    )
    object: str = Field(
        default=CHAT_COMPLETION_OBJECT,
        description="The object type, always 'chat.completion'.",
    )
    created: int = Field(
        default_factory=get_current_timestamp,
        description="The Unix timestamp when the completion was created.",
    )
    model: str = Field(description="The model used for the chat completion.")
    choices: list[ChatCompletionChoice] = Field(
        description="A list of chat completion choices."
    )
    usage: ChatCompletionUsage = Field(
        description="Usage statistics for the completion request."
    )
    system_fingerprint: str | None = Field(
        default=None,
        description="System fingerprint representing the backend configuration.",
    )


class ChatCompletionChunk(BaseModel):
    """A chunk of a chat completion."""

    id: str = Field(description="A unique identifier for the chat completion.")
    object: str = Field(
        default=CHAT_COMPLETION_CHUNK_OBJECT,
        description="The object type, always 'chat.completion.chunk'.",
    )
    created: int = Field(
        description="The Unix timestamp when the completion was created."
    )
    model: str = Field(description="The model used for the chat completion.")
    choices: list[ChatCompletionChunkChoice] = Field(
        description="A list of chat completion choices."
    )
    usage: ChatCompletionUsage | None = Field(
        default=None,
        description="Usage statistics for the completion request.",
    )
    system_fingerprint: str | None = Field(
        default=None,
        description="System fingerprint representing the backend configuration.",
    )
