from __future__ import annotations

import secrets
import time
from enum import Enum
from typing import Literal

from pydantic import BaseModel, Field

from orchard.server.models.responses.format import ResponseFormat
from orchard.server.models.responses.tools import (
    Function,
    FunctionID,
    ToolUseMode,
)

# --- Constants ---
RESPONSE_ID_PREFIX = "resp_"
RESPONSE_OBJECT = "response"
MESSAGE_ID_PREFIX = "msg_"
MESSAGE_OBJECT = "message"
OUTPUT_TEXT_OBJECT = "output_text"
FUNCTION_CALL_ID_PREFIX = "fc_"
TOOL_CALL_ID_PREFIX = "call_"


def generate_message_id(prefix: str = MESSAGE_ID_PREFIX) -> str:
    random_part = secrets.token_urlsafe(22)
    return f"{prefix}{random_part}"


def generate_response_id(prefix: str = RESPONSE_ID_PREFIX) -> str:
    random_part = secrets.token_urlsafe(22)  # Adjust length as needed
    return f"{prefix}{random_part}"


def generate_function_call_id(prefix: str = FUNCTION_CALL_ID_PREFIX) -> str:
    random_part = secrets.token_urlsafe(22)
    return f"{prefix}{random_part}"


def generate_tool_call_id(prefix: str = TOOL_CALL_ID_PREFIX) -> str:
    random_part = secrets.token_urlsafe(22)
    return f"{prefix}{random_part}"


def get_current_timestamp() -> int:
    return int(time.time())


class OutputTextContent(BaseModel):
    """Represents the text content of an output message."""

    type: Literal["output_text"] = OUTPUT_TEXT_OBJECT
    text: str = Field(description="The generated text content.")


class OutputStatus(Enum):
    """Represents the status of a message."""

    COMPLETED = "completed"
    INCOMPLETE = "incomplete"
    IN_PROGRESS = "in_progress"
    FAILED = "failed"


class OutputMessageType(Enum):
    """Represents the type of a message."""

    MESSAGE = "message"
    FUNCTION_CALL = "function_call"
    REASONING = "reasoning"


class OutputMessage(BaseModel):
    """Represents the message output within a response."""

    type: OutputMessageType = Field(
        default=OutputMessageType.MESSAGE, description="The type of the message."
    )
    status: OutputStatus = Field(
        default=OutputStatus.COMPLETED, description="The status of the message."
    )
    role: Literal["assistant"] = "assistant"
    id: str = Field(
        default_factory=generate_message_id, description="Unique message identifier."
    )
    content: list[OutputTextContent | OutputFunctionCall] = Field(
        description="Content generated by the assistant."
    )


class OutputFunctionCall(BaseModel):
    """Represents the function call output within a response."""

    type: Literal["function_call"] = "function_call"
    name: str = Field(description="The name of the function to call.")
    id: str = Field(
        default_factory=generate_function_call_id,
        description="Unique function call identifier.",
    )
    call_id: str = Field(
        default_factory=generate_tool_call_id,
        description="Unique tool call identifier.",
    )
    arguments: str = Field(description="The arguments to pass to the function.")
    status: OutputStatus = Field(
        default=OutputStatus.COMPLETED, description="The status of the function call."
    )


class ResponseUsage(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`."
    )


class ResponseObject(BaseModel):
    """Defines the response schema for the /v1/responses endpoint."""

    id: str = Field(
        default_factory=generate_response_id,
        description="Unique identifier for the response.",
    )
    object: Literal["response"] = RESPONSE_OBJECT
    created_at: int = Field(
        default_factory=get_current_timestamp,
        description="Unix timestamp of response creation.",
    )
    parallel_tool_calls: bool = Field(
        default=False,
        description="Whether to allow the model to run tool calls in parallel.",
    )
    temperature: float | None = Field(
        default=None,
        description="Sampling temperature.",
    )
    top_p: float | None = Field(
        default=None,
        description="Nucleus sampling threshold.",
    )
    top_k: int | None = Field(
        default=None,
        description="Controls the number of tokens considered at each step.",
    )
    min_p: float | None = Field(
        default=None,
        description="Minimum probability threshold for token consideration.",
    )
    tool_choice: ToolUseMode | FunctionID = Field(
        default=ToolUseMode.AUTO,
        description="How the model should select which tool (or tools) to use when generating a response.",
    )
    tools: list[Function] = Field(
        default_factory=list,
        description="A list of tools that the model can use to generate a response.",
    )
    status: OutputStatus = OutputStatus.COMPLETED
    model: str = Field(description="Model ID used for the response.")
    output: list[OutputMessage | OutputFunctionCall] = Field(
        description="List containing the output message(s)."
    )
    usage: ResponseUsage = Field(
        description="Usage statistics for the response request."
    )
    text: ResponseFormat | None = Field(
        default=None,
        description="The response format to use.",
    )
