"""Anthropic Messages API request/response schemas.

Reference: https://platform.claude.com/docs/en/api/messages
"""

from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, model_validator

# --- Content Block Types ---


class ImageSource(BaseModel):
    """Image source for image content blocks."""

    type: Literal["base64"] = "base64"
    media_type: Literal["image/jpeg", "image/png", "image/gif", "image/webp"]
    data: str  # Base64-encoded image data


class TextBlockParam(BaseModel):
    """Text content block in a message."""

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


class ImageBlockParam(BaseModel):
    """Image content block in a message."""

    type: Literal["image"] = "image"
    source: ImageSource


class ToolUseBlockParam(BaseModel):
    """Tool use content block in a message (assistant requesting tool execution)."""

    type: Literal["tool_use"] = "tool_use"
    id: str
    name: str
    input: dict[str, Any]


class ToolResultBlockParam(BaseModel):
    """Tool result content block in a message (user providing tool output)."""

    type: Literal["tool_result"] = "tool_result"
    tool_use_id: str
    content: str | list[TextBlockParam]


class AnthropicToolDefinition(BaseModel):
    """Tool definition for Anthropic Messages API."""

    name: str
    description: str
    input_schema: dict[str, Any]


# Union type for content blocks
ContentBlock = TextBlockParam | ImageBlockParam | ToolUseBlockParam | ToolResultBlockParam


# --- Message Types ---


class MessageParam(BaseModel):
    """A message in the conversation.

    Note: System messages are NOT allowed here - use the system parameter instead.
    """

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

    @model_validator(mode="after")
    def validate_content_length(self) -> "MessageParam":
        if isinstance(self.content, list) and len(self.content) > 256:
            raise ValueError("content may contain at most 256 blocks")
        return self


# --- Request Model ---


class AnthropicMessagesRequest(BaseModel):
    """Anthropic Messages API request.

    Reference: https://platform.claude.com/docs/en/api/messages

    Key differences from OpenAI:
    - max_tokens is REQUIRED (not optional)
    - system is a separate field (not in messages array)
    - content can be string or array of content blocks
    - temperature range is 0.0 to 1.0 (not 0.0 to 2.0)
    """

    model_config = ConfigDict(
        json_schema_extra={
            "examples": [
                {
                    "model": "mlx-community/Qwen3-0.6B-4bit-DWQ",
                    "max_tokens": 256,
                    "messages": [{"role": "user", "content": "What is the capital of France?"}],
                    "system": "You are a helpful assistant.",
                    "temperature": 0.7,
                    "stream": False,
                }
            ]
        }
    )

    model: str
    max_tokens: int = Field(ge=1, le=128000)  # Required - no default
    messages: list[MessageParam] = Field(..., max_length=1024)
    system: str | list[TextBlockParam] | None = None
    temperature: float = Field(default=1.0, ge=0.0, le=1.0)
    top_p: float | None = None
    top_k: int | None = None
    stop_sequences: list[str] | None = Field(default=None, max_length=16)
    stream: bool = False
    metadata: dict[str, Any] | None = None
    tools: list[AnthropicToolDefinition] | None = Field(default=None, max_length=256)
    tool_choice: dict[str, Any] | None = None

    @model_validator(mode="after")
    def validate_system_length(self) -> "AnthropicMessagesRequest":
        if isinstance(self.system, list) and len(self.system) > 64:
            raise ValueError("system may contain at most 64 blocks")
        return self


# --- Response Models ---


class ThinkingBlock(BaseModel):
    """Thinking (reasoning) content block in response."""

    type: Literal["thinking"] = "thinking"
    thinking: str


class TextBlock(BaseModel):
    """Text content block in response."""

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


class ToolUseBlock(BaseModel):
    """Tool use content block in response."""

    type: Literal["tool_use"] = "tool_use"
    id: str
    name: str
    input: dict[str, Any]


class Usage(BaseModel):
    """Token usage statistics."""

    input_tokens: int
    output_tokens: int


class AnthropicMessagesResponse(BaseModel):
    """Anthropic Messages API response.

    Reference: https://platform.claude.com/docs/en/api/messages
    """

    id: str
    type: Literal["message"] = "message"
    role: Literal["assistant"] = "assistant"
    model: str
    content: list[ThinkingBlock | TextBlock | ToolUseBlock]
    stop_reason: Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"] | None
    stop_sequence: str | None = None
    usage: Usage


# --- Streaming Event Models ---


class MessageStartMessage(BaseModel):
    """Message object in message_start event."""

    id: str
    type: Literal["message"] = "message"
    role: Literal["assistant"] = "assistant"
    content: list[Any] = Field(default_factory=list)
    model: str
    stop_reason: None = None
    stop_sequence: None = None
    usage: Usage


class MessageStartEvent(BaseModel):
    """SSE event: message_start."""

    type: Literal["message_start"] = "message_start"
    message: MessageStartMessage


class ContentBlockStartBlock(BaseModel):
    """Content block in content_block_start event."""

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


class ThinkingStartBlock(BaseModel):
    """Thinking block in content_block_start event."""

    type: Literal["thinking"] = "thinking"
    thinking: str = ""


class ToolUseStartBlock(BaseModel):
    """Tool use block in content_block_start event."""

    type: Literal["tool_use"] = "tool_use"
    id: str
    name: str
    input: dict[str, Any] = Field(default_factory=dict)


class ContentBlockStartEvent(BaseModel):
    """SSE event: content_block_start."""

    type: Literal["content_block_start"] = "content_block_start"
    index: int
    content_block: ThinkingStartBlock | ContentBlockStartBlock | ToolUseStartBlock


class TextDelta(BaseModel):
    """Text delta in content_block_delta event."""

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


class ThinkingDelta(BaseModel):
    """Thinking delta in content_block_delta event."""

    type: Literal["thinking_delta"] = "thinking_delta"
    thinking: str


class InputJsonDelta(BaseModel):
    """Input JSON delta in content_block_delta event for tool use."""

    type: Literal["input_json_delta"] = "input_json_delta"
    partial_json: str


class ContentBlockDeltaEvent(BaseModel):
    """SSE event: content_block_delta."""

    type: Literal["content_block_delta"] = "content_block_delta"
    index: int
    delta: ThinkingDelta | TextDelta | InputJsonDelta


class ContentBlockStopEvent(BaseModel):
    """SSE event: content_block_stop."""

    type: Literal["content_block_stop"] = "content_block_stop"
    index: int


class MessageDeltaDelta(BaseModel):
    """Delta in message_delta event."""

    stop_reason: Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"] | None
    stop_sequence: str | None = None


class MessageDeltaUsage(BaseModel):
    """Usage in message_delta event."""

    output_tokens: int


class MessageDeltaEvent(BaseModel):
    """SSE event: message_delta."""

    type: Literal["message_delta"] = "message_delta"
    delta: MessageDeltaDelta
    usage: MessageDeltaUsage


class MessageStopEvent(BaseModel):
    """SSE event: message_stop."""

    type: Literal["message_stop"] = "message_stop"


# --- Helper Functions ---


def extract_anthropic_content(
    content: str | list[ContentBlock],
) -> str:
    """Extract text from Anthropic content (string or content blocks).

    Args:
        content: Either a string or list of content blocks

    Returns:
        Concatenated text content
    """
    if isinstance(content, str):
        return content

    text_parts: list[str] = []
    for block in content:
        if isinstance(block, dict):
            # Handle dict form (from JSON parsing)
            if block.get("type") == "text":
                text_parts.append(block.get("text", ""))
        elif isinstance(block, TextBlockParam):
            # Handle Pydantic model form
            text_parts.append(block.text)

    return " ".join(text_parts)
