from __future__ import annotations

"""Anthropic provider implementation."""

import os
from typing import Any, AsyncIterator

import httpx

from agentrspy.core import (
    CompletionRequest,
    CompletionResponse,
    LlmProvider,
    Message,
    ProviderError,
    Role,
    StopReason,
    StreamChunk,
    ToolCall,
    Usage,
    map_stop_reason,
)


class AnthropicProvider(LlmProvider):
    """Anthropic provider."""

    def __init__(
        self, *, client: httpx.AsyncClient, api_key: str, base_url: str, model: str
    ) -> None:
        self.client = client
        self.api_key = api_key
        self.base_url = base_url
        self.model = model

    @classmethod
    def from_env(cls) -> "AnthropicBuilder":
        return AnthropicBuilder(api_key=os.getenv("ANTHROPIC_API_KEY"))

    @classmethod
    def builder(cls) -> "AnthropicBuilder":
        return AnthropicBuilder()

    async def complete(self, request: CompletionRequest) -> CompletionResponse:
        messages: list[dict[str, Any]] = []
        system = request.system or ""
        for message in request.messages:
            if message.role is Role.SYSTEM:
                if system:
                    system += "\n"
                system += message.text_content()
                continue
            messages.append(
                {
                    "role": "assistant" if message.role is Role.ASSISTANT else "user",
                    "content": [{"type": "text", "text": message.text_content()}],
                }
            )
        tools = None
        if request.tools is not None:
            tools = [
                {
                    "name": tool.name,
                    "description": tool.description,
                    "input_schema": tool.schema,
                }
                for tool in request.tools
            ]
        try:
            response = await self.client.post(
                f"{self.base_url.rstrip('/')}/messages",
                headers={"x-api-key": self.api_key, "anthropic-version": "2023-06-01"},
                json={
                    "model": request.model or self.model,
                    "messages": messages,
                    "system": system,
                    "max_tokens": request.max_tokens or 2048,
                    "temperature": request.temperature,
                    "tools": tools,
                },
            )
        except httpx.HTTPError as error:
            raise ProviderError.http(str(error)) from error
        if response.status_code >= 400:
            raise ProviderError.api(response.status_code, response.text)
        try:
            payload = response.json()
        except ValueError as error:
            raise ProviderError.parse(str(error)) from error
        text = ""
        tool_calls: list[ToolCall] = []
        content = payload.get("content")
        if isinstance(content, list):
            for item in content:
                if item.get("type") == "text" and isinstance(item.get("text"), str):
                    text += item["text"]
                elif item.get("type") == "tool_use":
                    tool_calls.append(
                        ToolCall(
                            id=str(item.get("id") or ""),
                            name=str(item.get("name") or ""),
                            arguments=item.get("input") or {},
                        )
                    )
        message = Message.assistant(text)
        if tool_calls:
            message.tool_calls = tool_calls
        usage_value = payload.get("usage") or {}
        usage = Usage(
            input_tokens=int(usage_value.get("input_tokens") or 0),
            output_tokens=int(usage_value.get("output_tokens") or 0),
            total_tokens=int(usage_value.get("input_tokens") or 0)
            + int(usage_value.get("output_tokens") or 0),
        )
        stop_reason = (
            map_stop_reason(payload.get("stop_reason", "stop"))
            if payload.get("stop_reason")
            else StopReason.STOP
        )
        return CompletionResponse(
            message=message,
            stop_reason=stop_reason,
            usage=usage,
            model=str(payload.get("model") or ""),
            raw=payload,
        )

    async def stream(self, request: CompletionRequest) -> AsyncIterator[StreamChunk]:
        del request
        raise ProviderError.unsupported(
            "Anthropic streaming is not yet implemented in this SDK scaffold"
        )

    def name(self) -> str:
        return "anthropic"


class AnthropicBuilder:
    """Builder for `AnthropicProvider`."""

    def __init__(
        self,
        *,
        api_key: str | None = None,
        base_url: str = "https://api.anthropic.com/v1",
        model: str = "claude-3-5-sonnet-latest",
        client: httpx.AsyncClient | None = None,
    ) -> None:
        self._api_key = api_key
        self._base_url = base_url
        self._model = model
        self._client = client or httpx.AsyncClient()

    def api_key(self, api_key: str) -> "AnthropicBuilder":
        self._api_key = api_key
        return self

    def base_url(self, base_url: str) -> "AnthropicBuilder":
        self._base_url = base_url
        return self

    def model(self, model: str) -> "AnthropicBuilder":
        self._model = model
        return self

    def build(self) -> AnthropicProvider:
        if not self._api_key:
            raise ProviderError.missing_api_key("ANTHROPIC_API_KEY")
        return AnthropicProvider(
            client=self._client,
            api_key=self._api_key,
            base_url=self._base_url,
            model=self._model,
        )
