"""Schemas for the LLM server."""

from collections.abc import Iterable
from typing import Literal

from mcp.types import Tool
from pydantic import BaseModel, Field


class TextBlock(BaseModel):
    """A schema for transporting text between the LLM and the MCP Client."""

    text: str = Field(description="The text content of the block.")

    type: Literal["text"] = "text"


class ToolUseBlock(BaseModel):
    """A schema for transporting tool requests between the LLM and the MCP Client."""

    id: str = Field(description="Unique identifier for the tool use request.")

    arguments: object = Field(description="Arguments for the tool use request.")

    name: str = Field(description="Name of the tool being used.")

    type: Literal["tool_use"] = "tool_use"


class ToolResultBlock(BaseModel):
    """A schema for transporting tool results between the LLM and the MCP Client."""

    tool_use_id: str = Field(description="Unique identifier for the tool use request.")

    name: str = Field(description="Name of the tool that was used.")

    content: str | Iterable[TextBlock] = Field(description="Content returned by the tool.")

    is_error: bool = Field(description="Indicates if the tool result is an error.")

    type: Literal["tool_result"] = "tool_result"


class Usage(BaseModel):
    """Generic usage tracking for any LLM provider."""

    input_tokens: int = Field(description="Number of input tokens")
    output_tokens: int = Field(description="Number of output tokens")
    cache_creation_input_tokens: int | None = Field(
        default=None, description="Number of tokens used for cache creation"
    )
    cache_read_input_tokens: int | None = Field(
        default=None, description="Number of tokens read from cache"
    )


Content = list[TextBlock | ToolUseBlock | ToolResultBlock]


class MessageBlock(BaseModel):
    """A message object for the request from the client."""

    content: Content = Field(
        description="Content of the message, which can include text and tool uses."
    )
    role: Literal["user", "assistant"]


class TextGenerationPayload(BaseModel):
    """The payload for the request."""

    messages: list[MessageBlock] = Field(description="Messages to be processed by the LLM.")
    tools: list[Tool] = Field(
        default_factory=list, description="Tools available for the LLM to use."
    )


class Message(BaseModel):
    """A message containing content and metadata."""

    id: str = Field(
        description="Unique object identifier.",
    )

    content: Content = Field(description="Content generated by the model.")

    model: str = Field(
        description="The model that completed the prompt.",
    )

    role: str = Field(
        default="assistant", description="Conversational role of the generated message."
    )

    stop_reason: str | None = Field(
        default=None,
        description="Reason for stopping generation, if applicable.",
    )

    usage: Usage | None = Field(default=None, description="Token usage information.")
