"""Anthropic API provider with streaming support."""

from __future__ import annotations

import json
from typing import Any, Generator

import anthropic

from tsumugi.providers.base import (
    AssistantMessage,
    ContentBlock,
    ProviderBase,
    StreamEvent,
    StreamEventType,
    TextBlock,
    ToolCallBlock,
    ToolDefinition,
    Usage,
)


class AnthropicProvider(ProviderBase):
    """Provider for Anthropic's Claude API."""

    def __init__(self, api_key: str, model: str = "claude-sonnet-4-20250514"):
        self.client = anthropic.Anthropic(api_key=api_key)
        self.model = model

    def _convert_tools(
        self, tools: list[ToolDefinition] | None
    ) -> list[dict[str, Any]] | anthropic.NotGiven:
        """Convert tool definitions to Anthropic format."""
        if not tools:
            return anthropic.NOT_GIVEN
        return [
            {
                "name": t.name,
                "description": t.description,
                "input_schema": t.parameters,
            }
            for t in tools
        ]

    def stream(
        self,
        messages: list[dict[str, Any]],
        system: str,
        tools: list[ToolDefinition] | None = None,
        max_tokens: int = 8192,
    ) -> Generator[StreamEvent, None, None]:
        """Stream a response from Claude."""
        with self.client.messages.stream(
            model=self.model,
            max_tokens=max_tokens,
            system=system,
            messages=messages,
            tools=self._convert_tools(tools),
        ) as stream:
            current_tool_id = ""
            current_tool_name = ""
            current_tool_input = ""

            for event in stream:
                if event.type == "content_block_start":
                    block = event.content_block
                    if block.type == "tool_use":
                        current_tool_id = block.id
                        current_tool_name = block.name
                        current_tool_input = ""
                        yield StreamEvent(
                            type=StreamEventType.TOOL_CALL_START,
                            tool_call_id=block.id,
                            tool_name=block.name,
                        )

                elif event.type == "content_block_delta":
                    delta = event.delta
                    if delta.type == "text_delta":
                        yield StreamEvent(
                            type=StreamEventType.TEXT_DELTA,
                            text=delta.text,
                        )
                    elif delta.type == "input_json_delta":
                        current_tool_input += delta.partial_json
                        yield StreamEvent(
                            type=StreamEventType.TOOL_CALL_DELTA,
                            tool_call_id=current_tool_id,
                            tool_name=current_tool_name,
                            tool_input_delta=delta.partial_json,
                        )

                elif event.type == "content_block_stop":
                    if current_tool_id:
                        yield StreamEvent(
                            type=StreamEventType.TOOL_CALL_END,
                            tool_call_id=current_tool_id,
                            tool_name=current_tool_name,
                            tool_input_delta=current_tool_input,
                        )
                        current_tool_id = ""
                        current_tool_name = ""
                        current_tool_input = ""

                elif event.type == "message_stop":
                    final = stream.get_final_message()
                    yield StreamEvent(
                        type=StreamEventType.MESSAGE_END,
                        stop_reason=final.stop_reason,
                        input_tokens=final.usage.input_tokens,
                        output_tokens=final.usage.output_tokens,
                    )

    def build_assistant_message(
        self, events: list[StreamEvent]
    ) -> AssistantMessage:
        """Reconstruct assistant message from stream events."""
        content: list[ContentBlock] = []
        text_parts: list[str] = []
        tool_inputs: dict[str, str] = {}  # tool_id -> accumulated JSON
        stop_reason = ""
        input_tokens = 0
        output_tokens = 0

        for ev in events:
            if ev.type == StreamEventType.TEXT_DELTA:
                text_parts.append(ev.text)

            elif ev.type == StreamEventType.TOOL_CALL_START:
                # Flush accumulated text before a tool call
                if text_parts:
                    content.append(TextBlock(text="".join(text_parts)))
                    text_parts = []
                tool_inputs[ev.tool_call_id] = ""

            elif ev.type == StreamEventType.TOOL_CALL_DELTA:
                tool_inputs[ev.tool_call_id] = (
                    tool_inputs.get(ev.tool_call_id, "") + ev.tool_input_delta
                )

            elif ev.type == StreamEventType.TOOL_CALL_END:
                raw = tool_inputs.get(ev.tool_call_id, ev.tool_input_delta)
                try:
                    parsed = json.loads(raw) if raw else {}
                except json.JSONDecodeError:
                    parsed = {}
                content.append(
                    ToolCallBlock(
                        id=ev.tool_call_id,
                        name=ev.tool_name,
                        input=parsed,
                    )
                )

            elif ev.type == StreamEventType.MESSAGE_END:
                stop_reason = ev.stop_reason
                input_tokens = ev.input_tokens
                output_tokens = ev.output_tokens

        # Flush remaining text
        if text_parts:
            content.append(TextBlock(text="".join(text_parts)))

        return AssistantMessage(
            content=content,
            stop_reason=stop_reason,
            usage=Usage(input_tokens=input_tokens, output_tokens=output_tokens),
        )
