"""Anthropic provider implementation."""

from __future__ import annotations

from typing import TYPE_CHECKING

from koan.llm.retry import llm_retry
from koan.types import LLMResponse, Message, StreamDelta, Usage

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from koan.llm.config import LLMConfig


class AnthropicProvider:
    """Async Anthropic messages provider."""

    def __init__(self, api_key: str) -> None:
        try:
            from anthropic import AsyncAnthropic
        except ImportError as e:
            raise ImportError(
                "Anthropic SDK not installed. Install with: pip install koan[anthropic]"
            ) from e
        self._client = AsyncAnthropic(api_key=api_key)

    def _build_kwargs(self, config: LLMConfig, messages: list[Message]) -> dict:
        """Build the kwargs dict for Anthropic messages.create()."""
        system_text = ""
        conv_messages: list[dict] = []
        for m in messages:
            if m.role == "system":
                system_text = m.content
            elif m.role == "tool" and m.tool_call_id:
                # Anthropic tool results go as user messages with tool_result blocks
                conv_messages.append({
                    "role": "user",
                    "content": [
                        {
                            "type": "tool_result",
                            "tool_use_id": m.tool_call_id,
                            "content": m.content,
                        }
                    ],
                })
            elif m.role == "assistant" and m.tool_calls:
                # Anthropic expects tool_use blocks in assistant content
                content: list[dict] = []
                if m.content:
                    content.append({"type": "text", "text": m.content})
                for tc in m.tool_calls:
                    content.append({
                        "type": "tool_use",
                        "id": tc["id"],
                        "name": tc["name"],
                        "input": tc["arguments"] if isinstance(tc["arguments"], dict) else {},
                    })
                conv_messages.append({"role": "assistant", "content": content})
            else:
                conv_messages.append({"role": m.role, "content": m.content})

        kwargs: dict = {
            "model": config.model,
            "max_tokens": config.max_tokens,
            "messages": conv_messages,
            "temperature": config.temperature,
            "top_p": config.top_p,
            "timeout": config.timeout,
        }
        if system_text:
            kwargs["system"] = system_text

        # Convert OpenAI-format tools to Anthropic format
        if config.tools:
            kwargs["tools"] = [
                {
                    "name": t["function"]["name"],
                    "description": t["function"].get("description", ""),
                    "input_schema": t["function"].get("parameters", {"type": "object", "properties": {}}),
                }
                for t in config.tools
            ]

        return kwargs

    @llm_retry
    async def complete(self, config: LLMConfig, messages: list[Message]) -> LLMResponse:
        """Call Anthropic Messages API and return a normalized response."""
        kwargs = self._build_kwargs(config, messages)
        response = await self._client.messages.create(**kwargs)

        text = ""
        tool_calls = []
        for block in response.content:
            if block.type == "text":
                text += block.text
            elif block.type == "tool_use":
                tool_calls.append({
                    "id": block.id,
                    "name": block.name,
                    "arguments": block.input,
                })

        raw: dict | None = {"tool_calls": tool_calls} if tool_calls else None

        return LLMResponse(
            text=text,
            model=response.model,
            provider="anthropic",
            usage=Usage(
                input_tokens=response.usage.input_tokens,
                output_tokens=response.usage.output_tokens,
            ),
            stop_reason=response.stop_reason,
            raw=raw,
        )

    async def stream(self, config: LLMConfig, messages: list[Message]) -> AsyncIterator[StreamDelta]:
        """Stream Anthropic Messages API and yield normalized deltas."""
        kwargs = self._build_kwargs(config, messages)

        async with self._client.messages.stream(**kwargs) as stream:
            async for event in stream:
                if event.type == "content_block_delta":
                    delta = event.delta
                    if delta.type == "text_delta":
                        yield StreamDelta(
                            text=delta.text,
                            provider="anthropic",
                        )
                    elif delta.type == "input_json_delta":
                        yield StreamDelta(
                            provider="anthropic",
                            tool_args_delta=delta.partial_json,
                        )
                elif event.type == "content_block_start":
                    block = event.content_block
                    if block.type == "tool_use":
                        yield StreamDelta(
                            provider="anthropic",
                            tool_call_id=block.id,
                            tool_name=block.name,
                        )
                elif event.type == "message_delta":
                    yield StreamDelta(
                        provider="anthropic",
                        finish_reason=event.delta.stop_reason,
                        usage=Usage(output_tokens=event.usage.output_tokens) if event.usage else None,
                    )
                elif event.type == "message_start":
                    msg = event.message
                    yield StreamDelta(
                        provider="anthropic",
                        model=msg.model,
                        usage=Usage(
                            input_tokens=msg.usage.input_tokens,
                            output_tokens=msg.usage.output_tokens,
                        ),
                    )

    async def close(self) -> None:
        await self._client.close()
