"""
Type definitions for claude-py SDK.

All types match the structure from Claude Code API captures.
"""

from __future__ import annotations

from collections.abc import Awaitable, Callable
from typing import Any, Literal

from pydantic import BaseModel

ToolHandler = Callable[[str, dict[str, Any]], Awaitable[str]]


class ToolUse(BaseModel):
    """Tool use request from assistant."""

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


class TextContent(BaseModel):
    """Text content block."""

    type: Literal["text"]
    text: str


class ToolResult(BaseModel):
    """Tool execution result sent back to assistant."""

    type: Literal["tool_result"]
    tool_use_id: str
    content: str
    is_error: bool | None = False


class CacheControl(BaseModel):
    """Cache control for prompt caching."""

    type: Literal["ephemeral"]


class SystemPromptBlock(BaseModel):
    """System prompt block with optional caching."""

    type: Literal["text"]
    text: str
    cache_control: CacheControl | None = None


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

    input_tokens: int
    output_tokens: int
    cache_creation_input_tokens: int | None = 0
    cache_read_input_tokens: int | None = 0


class AssistantMessage(BaseModel):
    """Complete assistant response."""

    id: str
    type: Literal["message"]
    role: Literal["assistant"]
    content: list[dict[str, Any]]
    model: str
    stop_reason: str | None = None
    stop_sequence: str | None = None
    usage: Usage | None = None

    @property
    def text(self) -> str:
        return "".join(b["text"] for b in self.content if b.get("type") == "text")

    @property
    def tool_calls(self) -> list[ToolUse]:
        return [ToolUse.model_validate(b) for b in self.content if b.get("type") == "tool_use"]

    @property
    def has_tool_use(self) -> bool:
        return self.stop_reason == "tool_use"


# Streaming event types
class StreamEvent(BaseModel):
    """Base class for streaming events."""

    type: str


class MessageStartEvent(StreamEvent):
    """Stream event: message start."""

    type: Literal["message_start"] = "message_start"
    message: dict[str, Any]


class ContentBlockStartEvent(StreamEvent):
    """Stream event: content block start."""

    type: Literal["content_block_start"] = "content_block_start"
    index: int
    content_block: dict[str, Any]


class ContentBlockDeltaEvent(StreamEvent):
    """Stream event: content block delta (incremental text)."""

    type: Literal["content_block_delta"] = "content_block_delta"
    index: int
    delta: dict[str, Any]


class ContentBlockStopEvent(StreamEvent):
    """Stream event: content block stop."""

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


class MessageDeltaEvent(StreamEvent):
    """Stream event: message delta (stop_reason, etc)."""

    type: Literal["message_delta"] = "message_delta"
    delta: dict[str, Any]
    usage: dict[str, Any] | None = None


class MessageStopEvent(StreamEvent):
    """Stream event: message stop."""

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


class PingEvent(StreamEvent):
    """Stream event: ping (keepalive)."""

    type: Literal["ping"] = "ping"


class ErrorEvent(StreamEvent):
    """Stream event: error."""

    type: Literal["error"] = "error"
    error: dict[str, Any]


class AgentOptions(BaseModel):
    """Configuration options for Claude agent."""

    model: str = "claude-sonnet-4-6"
    max_tokens: int = 32000
    stream: bool = True
    effort: Literal["low", "medium", "high"] | None = None
    allowed_tools: list[str] | None = None
    disallowed_tools: list[str] | None = None
    system_prompt: str | None = None
    append_system_prompt: str | None = None
    cwd: str = "/home/claude/host"
    session_id: str | None = None
    resume: str | None = None
    device_id: str | None = None
    account_uuid: str | None = None


class StreamChunk(BaseModel):
    """Represents a chunk of streaming data."""

    event_type: str
    data: dict[str, Any]
    text_delta: str | None = None  # Extracted text for convenience
    content_block: dict[str, Any] | None = None


class StreamResponse:
    """Wrapper for streaming response that accumulates chunks."""

    def __init__(self):
        self.chunks: list[StreamChunk] = []
        self.accumulated_text = ""
        self.content_blocks: list[dict[str, Any]] = []
        self.message_data: dict[str, Any] = {}
        self.stop_reason: str | None = None
        self.usage: dict[str, Any] | None = None
        self.current_block: dict[str, Any] | None = None

    def add_chunk(self, chunk: StreamChunk):
        """Add a chunk and update accumulated state."""
        self.chunks.append(chunk)

        if chunk.event_type == "message_start":
            self.message_data = chunk.data.get("message", {})

        elif chunk.event_type == "content_block_start":
            self.current_block = chunk.data.get("content_block", {})

        elif chunk.event_type == "content_block_delta":
            delta = chunk.data.get("delta", {})
            if delta.get("type") == "text_delta":
                text = delta.get("text", "")
                self.accumulated_text += text
                # Update current block with accumulated text
                if self.current_block:
                    self.current_block["text"] = self.accumulated_text

        elif chunk.event_type == "content_block_stop":
            if self.current_block:
                self.content_blocks.append(self.current_block)
                self.current_block = None
                self.accumulated_text = ""

        elif chunk.event_type == "message_delta":
            delta = chunk.data.get("delta", {})
            self.stop_reason = delta.get("stop_reason")
            if chunk.data.get("usage"):
                self.usage = chunk.data["usage"]

    def to_message(self) -> AssistantMessage:
        """Convert accumulated stream to a complete AssistantMessage."""
        return AssistantMessage(
            id=self.message_data.get("id", ""),
            type="message",
            role="assistant",
            content=self.content_blocks,
            model=self.message_data.get("model", ""),
            stop_reason=self.stop_reason,
            usage=Usage(**self.usage) if self.usage else None,
        )