"""Anthropic Claude Messages API client."""

from __future__ import annotations

import json
from collections.abc import AsyncIterator
from typing import Any

import httpx

from llm_rotator._types import LLMResponse, StreamChunk, ToolCall, Usage
from llm_rotator.clients import AbstractLLMClient
from llm_rotator.exceptions import KeyDeadError, ModelRateLimitError, ServerError

_BASE_URL = "https://api.anthropic.com"
_ANTHROPIC_VERSION = "2023-06-01"
_DEFAULT_MAX_TOKENS = 4096


class AnthropicClient(AbstractLLMClient):
    """Client for Anthropic Claude Messages API."""

    def __init__(self, timeout: float = 60.0) -> None:
        self._timeout = timeout

    async def generate(
        self,
        messages: list[dict],
        model: str,
        api_key: str,
        base_url: str | None = None,
        **kwargs: Any,
    ) -> LLMResponse:
        url = f"{base_url or _BASE_URL}/v1/messages"
        headers = _build_headers(api_key)
        body = _build_body(messages, model, **kwargs)

        async with httpx.AsyncClient(timeout=self._timeout) as client:
            resp = await client.post(url, json=body, headers=headers)

        _check_status(resp, model=model, api_key=api_key)

        data = resp.json()
        return _parse_response(data, model=model, api_key=api_key)

    async def stream(
        self,
        messages: list[dict],
        model: str,
        api_key: str,
        base_url: str | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[StreamChunk]:
        url = f"{base_url or _BASE_URL}/v1/messages"
        headers = _build_headers(api_key)
        body = _build_body(messages, model, **kwargs)
        body["stream"] = True

        # Track tool_use blocks being streamed
        pending_tool_calls: dict[int, dict] = {}  # index → {id, name, arguments}

        async with (
            httpx.AsyncClient(timeout=self._timeout) as client,
            client.stream("POST", url, json=body, headers=headers) as resp,
        ):
            if resp.status_code != 200:
                await resp.aread()
                _check_status(resp, model=model, api_key=api_key)

            prompt_tokens = 0
            completion_tokens = 0

            async for line in resp.aiter_lines():
                if not line.startswith("data: "):
                    continue
                payload = line[6:].strip()
                if not payload:
                    continue

                data = json.loads(payload)
                event_type = data.get("type", "")

                if event_type == "content_block_start":
                    block = data.get("content_block", {})
                    idx = data.get("index", 0)
                    if block.get("type") == "tool_use":
                        pending_tool_calls[idx] = {
                            "id": block.get("id", ""),
                            "name": block.get("name", ""),
                            "arguments": "",
                        }

                elif event_type == "content_block_delta":
                    delta_data = data.get("delta", {})
                    idx = data.get("index", 0)
                    if delta_data.get("type") == "input_json_delta":
                        if idx in pending_tool_calls:
                            pending_tool_calls[idx]["arguments"] += delta_data.get(
                                "partial_json", ""
                            )
                    else:
                        text = delta_data.get("text", "")
                        if text:
                            yield StreamChunk(delta=text)

                elif event_type == "message_start":
                    msg = data.get("message", {})
                    msg_usage = msg.get("usage", {})
                    prompt_tokens = msg_usage.get("input_tokens", 0)

                elif event_type == "message_delta":
                    usage_data = data.get("usage", {})
                    completion_tokens = usage_data.get("output_tokens", 0)

                elif event_type == "message_stop":
                    usage = Usage(
                        prompt_tokens=prompt_tokens,
                        completion_tokens=completion_tokens,
                        total_tokens=prompt_tokens + completion_tokens,
                    )
                    final_tool_calls = _finalize_tool_calls(pending_tool_calls)
                    yield StreamChunk(
                        delta="", usage=usage, done=True, tool_calls=final_tool_calls
                    )
                    return


def _build_headers(api_key: str) -> dict[str, str]:
    return {
        "x-api-key": api_key,
        "anthropic-version": _ANTHROPIC_VERSION,
        "content-type": "application/json",
    }


def _build_body(messages: list[dict], model: str, **kwargs: Any) -> dict:
    """Build Anthropic request body from OpenAI-format messages."""
    system_parts: list[str] = []
    translated: list[dict] = []

    for msg in messages:
        role = msg["role"]
        content = msg["content"]
        if role == "system":
            system_parts.append(content)
        else:
            translated.append({"role": role, "content": content})

    body: dict = {
        "model": model,
        "messages": translated,
        "max_tokens": _DEFAULT_MAX_TOKENS,
    }

    # Tools: translate OpenAI format → Anthropic format
    tools = kwargs.get("tools")
    if tools:
        body["tools"] = [
            {
                "name": t.get("function", t)["name"],
                "description": t.get("function", t).get("description", ""),
                "input_schema": t.get("function", t).get("parameters", {}),
            }
            for t in tools
        ]

    # Tool choice
    tool_choice = kwargs.get("tool_choice")
    if tool_choice is not None:
        if tool_choice == "auto":
            body["tool_choice"] = {"type": "auto"}
        elif tool_choice == "none":
            # Anthropic doesn't support "none" directly; omit tools
            body.pop("tools", None)
        elif tool_choice == "required":
            body["tool_choice"] = {"type": "any"}
        elif isinstance(tool_choice, dict) and tool_choice.get("function"):
            body["tool_choice"] = {
                "type": "tool",
                "name": tool_choice["function"]["name"],
            }

    # JSON mode / structured output via system prompt
    response_format = kwargs.get("response_format")
    if response_format is not None:
        instruction = ""
        if isinstance(response_format, dict):
            if response_format.get("type") == "json_object":
                instruction = "You must respond with valid JSON only."
        elif hasattr(response_format, "model_json_schema"):
            schema = response_format.model_json_schema()
            instruction = (
                f"You must respond with valid JSON matching this schema: "
                f"{json.dumps(schema)}"
            )
        if instruction:
            system_parts.append(instruction)

    if system_parts:
        body["system"] = "\n\n".join(system_parts)

    return body


def _parse_response(data: dict, *, model: str, api_key: str) -> LLMResponse:
    """Parse Anthropic Messages API response."""
    content = ""
    tool_calls: list[ToolCall] = []

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

    usage_data = data.get("usage", {})
    input_tokens = usage_data.get("input_tokens", 0)
    output_tokens = usage_data.get("output_tokens", 0)

    return LLMResponse(
        content=content or None if tool_calls else content,
        usage=Usage(
            prompt_tokens=input_tokens,
            completion_tokens=output_tokens,
            total_tokens=input_tokens + output_tokens,
        ),
        model=data.get("model", model),
        provider="anthropic",
        key_alias=api_key[:8] + "...",
        tool_calls=tool_calls or None,
        raw=data,
    )


def _finalize_tool_calls(
    pending: dict[int, dict],
) -> list[ToolCall] | None:
    """Convert accumulated stream tool call fragments into ToolCall list."""
    if not pending:
        return None
    return [
        ToolCall(
            id=pending[idx]["id"],
            name=pending[idx]["name"],
            arguments=pending[idx]["arguments"],
        )
        for idx in sorted(pending)
    ]


def _check_status(
    resp: httpx.Response, *, model: str, api_key: str
) -> None:
    """Raise typed exceptions based on HTTP status code."""
    if resp.status_code < 400:
        return

    alias = api_key[:8] + "..."
    status = resp.status_code

    if status in (401, 403, 402):
        raise KeyDeadError(key_alias=alias, status_code=status)
    if status == 429:
        retry_after_raw = resp.headers.get("retry-after")
        retry_after = int(retry_after_raw) if retry_after_raw else None
        raise ModelRateLimitError(
            key_alias=alias, model=model, retry_after=retry_after
        )
    if status >= 500 or status == 529:
        raise ServerError(provider="anthropic", status_code=status)

    raise ServerError(
        provider="anthropic",
        status_code=status,
        message=f"Unexpected HTTP {status}",
    )
