"""Anthropic (Claude) provider for Logos ADK."""

from __future__ import annotations

import os
import httpx
from typing import Any, AsyncIterator

from logos_adk.errors import AuthenticationError, ProviderError
from logos_adk.providers import Provider, provider_registry
from logos_adk.types import (
    Message,
    Response,
    TextChunk,
    ToolCall,
    Usage,
)


def _messages_to_anthropic(messages: list[Message]) -> tuple[str | None, list[dict]]:
    """Convert Logos messages to Anthropic API format.

    Returns (system_prompt, messages_list).
    """
    system: str | None = None
    api_messages: list[dict[str, Any]] = []

    for msg in messages:
        if msg.role == "system":
            system = msg.content
            continue

        if msg.role == "tool":
            # Tool result message
            if msg.tool_result is not None:
                api_messages.append(
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "tool_result",
                                "tool_use_id": msg.tool_result.call_id,
                                "content": msg.tool_result.error or str(msg.tool_result.content),
                                "is_error": msg.tool_result.error is not None,
                            }
                        ],
                    }
                )
            continue

        if msg.role == "assistant" and msg.tool_calls:
            content: list[dict[str, Any]] = []
            if msg.content:
                content.append({"type": "text", "text": msg.content})
            for tc in msg.tool_calls:
                content.append(
                    {
                        "type": "tool_use",
                        "id": tc.id,
                        "name": tc.name,
                        "input": tc.arguments,
                    }
                )
            api_messages.append({"role": "assistant", "content": content})
            continue

        api_messages.append({"role": msg.role, "content": msg.content})

    return system, api_messages


def _tools_to_anthropic(tools: list[dict] | None) -> list[dict] | None:
    """Convert Logos tool schemas to Anthropic format."""
    if not tools:
        return None
    return [
        {
            "name": t["name"],
            "description": t.get("description", ""),
            "input_schema": t.get("parameters", {"type": "object", "properties": {}}),
        }
        for t in tools
    ]


def _parse_response(raw: Any) -> Response:
    """Parse an Anthropic response into a Logos Response."""
    content_parts: list[str] = []
    tool_calls: list[ToolCall] = []

    for block in raw.content:
        if block.type == "text":
            content_parts.append(block.text)
        elif block.type == "tool_use":
            tool_calls.append(ToolCall(id=block.id, name=block.name, arguments=block.input))

    if hasattr(raw, "usage"):
        input_tokens = raw.usage.input_tokens
        output_tokens = raw.usage.output_tokens
    else:
        input_tokens = 0
        output_tokens = 0

    return Response(
        content="".join(content_parts),
        tool_calls=tool_calls,
        usage=Usage(
            input_tokens=input_tokens,
            output_tokens=output_tokens,
        ),
        raw={"id": raw.id, "model": raw.model, "stop_reason": raw.stop_reason},
    )


class AnthropicProvider(Provider):
    """Provider for Anthropic Claude models."""

    _http_client_pool: dict[str | None, httpx.AsyncClient] = {}
    _client_pool: dict[tuple[str, str | None], Any] = {}

    @classmethod
    def _shared_http_client(cls, proxy: str | None) -> httpx.AsyncClient:
        if proxy not in cls._http_client_pool:
            kwargs: dict[str, Any] = {"timeout": 120.0}
            if proxy:
                kwargs["proxy"] = proxy
            cls._http_client_pool[proxy] = httpx.AsyncClient(**kwargs)
        return cls._http_client_pool[proxy]

    def __init__(
        self,
        model: str = "claude-sonnet-4-20250514",
        *,
        api_key: str | None = None,
        max_tokens: int = 4096,
        proxy: str | None = None,
    ) -> None:
        try:
            import anthropic
        except ImportError:
            raise ImportError(
                "anthropic package required. Install with: pip install logos-adk[anthropic]"
            )

        resolved_key = api_key or os.environ.get("ANTHROPIC_API_KEY")
        if not resolved_key:
            raise AuthenticationError(
                "Anthropic API key not found. Set ANTHROPIC_API_KEY env var "
                "or pass api_key= to the provider."
            )

        client_key = (resolved_key, proxy)
        if client_key not in self._client_pool:
            self._client_pool[client_key] = anthropic.AsyncAnthropic(
                api_key=resolved_key,
                http_client=self._shared_http_client(proxy),
            )
        self._client = self._client_pool[client_key]
        self._model = model
        self._max_tokens = max_tokens
        self.proxy = proxy

    async def generate(
        self,
        messages: list[Message],
        tools: list[dict] | None = None,
        output_schema: Any = None,
        **kwargs: Any,
    ) -> Response:
        system, api_messages = _messages_to_anthropic(messages)
        api_tools = _tools_to_anthropic(tools)

        create_kwargs: dict[str, Any] = {
            "model": self._model,
            "max_tokens": kwargs.pop("max_tokens", self._max_tokens),
            "messages": api_messages,
        }
        if system:
            create_kwargs["system"] = system
        if api_tools:
            create_kwargs["tools"] = api_tools
        create_kwargs.update(kwargs)

        try:
            raw = await self._client.messages.create(**create_kwargs)
        except Exception as e:
            raise ProviderError(f"Anthropic API error: {e}") from e

        return _parse_response(raw)

    async def stream(
        self,
        messages: list[Message],
        tools: list[dict] | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[Any]:
        system, api_messages = _messages_to_anthropic(messages)
        api_tools = _tools_to_anthropic(tools)

        create_kwargs: dict[str, Any] = {
            "model": self._model,
            "max_tokens": kwargs.pop("max_tokens", self._max_tokens),
            "messages": api_messages,
        }
        if system:
            create_kwargs["system"] = system
        if api_tools:
            create_kwargs["tools"] = api_tools
        create_kwargs.update(kwargs)

        try:
            async with self._client.messages.stream(**create_kwargs) as stream:
                async for text in stream.text_stream:
                    yield TextChunk(text=text)

                response = await stream.get_final_message()
                yield _parse_response(response)
        except Exception as e:
            raise ProviderError(f"Anthropic streaming error: {e}") from e


def _anthropic_factory(model_name: str) -> AnthropicProvider:
    return AnthropicProvider(model=model_name)


# Auto-register for claude-* model patterns
provider_registry.register("claude-*", _anthropic_factory)
