from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field


class AnthropicTextBlock(BaseModel):
    type: Literal["text"]
    text: str


class AnthropicImageSource(BaseModel):
    type: Literal["base64"]
    media_type: str
    data: str


class AnthropicImageBlock(BaseModel):
    type: Literal["image"]
    source: AnthropicImageSource


class AnthropicToolResultBlock(BaseModel):
    type: Literal["tool_result"]
    tool_use_id: str
    content: str | list[dict[str, Any]]
    is_error: bool | None = None


class AnthropicToolUseBlock(BaseModel):
    type: Literal["tool_use"]
    id: str
    name: str
    input: dict[str, Any]


class AnthropicThinkingBlock(BaseModel):
    type: Literal["thinking"]
    thinking: str


AnthropicContentBlock = (
    AnthropicTextBlock
    | AnthropicImageBlock
    | AnthropicToolResultBlock
    | AnthropicToolUseBlock
    | AnthropicThinkingBlock
)


class AnthropicMessage(BaseModel):
    role: Literal["user", "assistant"]
    content: str | list[AnthropicContentBlock]


class AnthropicTool(BaseModel):
    name: str
    description: str | None = None
    input_schema: dict[str, Any]


class AnthropicToolChoice(BaseModel):
    type: Literal["auto", "any", "tool", "none"]
    name: str | None = None


class AnthropicMetadata(BaseModel):
    user_id: str | None = None


class AnthropicMessagesRequest(BaseModel):
    model_config = ConfigDict(extra="allow")

    model: str
    messages: list[AnthropicMessage]
    max_tokens: int = Field(default=1024, ge=1)
    system: str | list[AnthropicTextBlock] | None = None
    metadata: AnthropicMetadata | None = None
    stop_sequences: list[str] | None = None
    stream: bool | None = False
    temperature: float | None = None
    top_p: float | None = None
    top_k: int | None = None
    tools: list[AnthropicTool] | None = None
    tool_choice: AnthropicToolChoice | None = None


class AnthropicUsage(BaseModel):
    input_tokens: int = 0
    output_tokens: int = 0
    cache_creation_input_tokens: int | None = None
    cache_read_input_tokens: int | None = None


class AnthropicMessageResponse(BaseModel):
    id: str
    type: Literal["message"] = "message"
    role: Literal["assistant"] = "assistant"
    content: list[AnthropicTextBlock | AnthropicToolUseBlock]
    model: str
    stop_reason: str | None
    stop_sequence: str | None = None
    usage: AnthropicUsage


class CountTokensRequest(BaseModel):
    model_config = ConfigDict(extra="allow")

    model: str
    messages: list[AnthropicMessage]
    system: str | list[AnthropicTextBlock] | None = None
    tools: list[AnthropicTool] | None = None


class CountTokensResponse(BaseModel):
    input_tokens: int


SUPPORTED_MODELS = [
    {
        "id": "claude-sonnet-4.6",
        "display_name": "Claude Sonnet 4.6",
        "owned_by": "anthropic",
    },
    {
        "id": "claude-opus-4.6",
        "display_name": "Claude Opus 4.6",
        "owned_by": "anthropic",
    },
    {
        "id": "claude-haiku-4.5",
        "display_name": "Claude Haiku 4.5",
        "owned_by": "anthropic",
    },
]