"""Anthropic Claude LLM provider."""

from __future__ import annotations

from collections.abc import AsyncIterator

try:
    import anthropic
except ModuleNotFoundError:
    anthropic = None  # type: ignore[assignment]

from agntspace.llm.base import LLMProvider, LLMResponse, Message, StreamChunk, ToolCall


class AnthropicProvider(LLMProvider):
    """LLM provider using the Anthropic Messages API."""

    def __init__(self, api_key: str, model: str = "claude-sonnet-4-20250514") -> None:
        self._client = anthropic.AsyncAnthropic(api_key=api_key)
        self._model = model

    async def complete(
        self,
        messages: list[Message],
        system: str | None = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
    ) -> LLMResponse:
        kwargs: dict = {
            "model": self._model,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "messages": [{"role": m.role, "content": m.content} for m in messages],
        }
        if system:
            kwargs["system"] = system

        resp = await self._client.messages.create(**kwargs)
        return LLMResponse(
            content=resp.content[0].text if resp.content else "",
            model=resp.model,
            input_tokens=resp.usage.input_tokens,
            output_tokens=resp.usage.output_tokens,
            cache_read_tokens=getattr(resp.usage, "cache_read_input_tokens", 0) or 0,
            cache_write_tokens=getattr(resp.usage, "cache_creation_input_tokens", 0) or 0,
            stop_reason=resp.stop_reason,
        )

    async def complete_with_tools(
        self,
        messages: list[Message],
        tools: list[dict],
        system: str | None = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
    ) -> LLMResponse:
        """Anthropic native tool use."""
        msg_dicts = [{"role": m.role, "content": m.content} for m in messages]

        kwargs: dict = {
            "model": self._model,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "messages": msg_dicts,
            "tools": tools,
        }
        if system:
            kwargs["system"] = system

        resp = await self._client.messages.create(**kwargs)

        # Extract text and tool calls
        text = ""
        tool_calls = []
        for block in resp.content:
            if hasattr(block, "text"):
                text += block.text
            elif block.type == "tool_use":
                tool_calls.append(ToolCall(id=block.id, name=block.name, input=block.input))

        return LLMResponse(
            content=text,
            model=resp.model,
            input_tokens=resp.usage.input_tokens,
            output_tokens=resp.usage.output_tokens,
            cache_read_tokens=getattr(resp.usage, "cache_read_input_tokens", 0) or 0,
            cache_write_tokens=getattr(resp.usage, "cache_creation_input_tokens", 0) or 0,
            stop_reason=resp.stop_reason,
            tool_calls=tool_calls,
            raw_content=resp.content,
        )

    async def stream(
        self,
        messages: list[Message],
        system: str | None = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
    ) -> AsyncIterator[StreamChunk]:
        kwargs: dict = {
            "model": self._model,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "messages": [{"role": m.role, "content": m.content} for m in messages],
        }
        if system:
            kwargs["system"] = system

        async with self._client.messages.stream(**kwargs) as stream:
            async for text in stream.text_stream:
                yield StreamChunk(delta=text)

            response = await stream.get_final_message()
            yield StreamChunk(
                delta="",
                done=True,
                response=LLMResponse(
                    content=response.content[0].text if response.content else "",
                    model=response.model,
                    input_tokens=response.usage.input_tokens,
                    output_tokens=response.usage.output_tokens,
                    cache_read_tokens=getattr(response.usage, "cache_read_input_tokens", 0) or 0,
                    cache_write_tokens=getattr(response.usage, "cache_creation_input_tokens", 0) or 0,
                    stop_reason=response.stop_reason,
                ),
            )
