"""Anthropic OAuth provider yielding LiteLLM-style ProviderChunk deltas."""

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

import anthropic

from ..base import OAuthProvider, OAuthProviderError, ProviderChunk
from ..registry import register_provider

_BETA_FEATURES = [
    "claude-code-20250219",
    "oauth-2025-04-20",
    "fine-grained-tool-streaming-2025-05-14",
    "interleaved-thinking-2025-05-14",
]
_DEFAULT_MAX_TOKENS = 4096
_MAX_RETRIES = 3
_TIMEOUT_SECONDS = 60.0
_THINKING_BUDGETS = {
    "low": 1024,
    "medium": 4096,
    "high": 8192,
    "xhigh": 16384,
}


def _sanitize_text(text: str) -> str:
    return text.encode("utf-8", errors="surrogatepass").decode("utf-8", errors="replace")


def _coerce_int(value: Any) -> int | None:
    if value is None:
        return None
    try:
        return int(value)
    except (TypeError, ValueError):
        return None


def _normalize_reasoning_effort(effort: str | None) -> str:
    if effort is None:
        return "none"
    normalized = effort.strip().lower()
    if normalized in {"none", "off", "disabled"}:
        return "none"
    if normalized in {"on", "enabled", "med"}:
        return "medium"
    if normalized in {"low", "minimal"}:
        return "low"
    if normalized in {"medium", "high", "xhigh"}:
        return normalized
    return "medium"


@register_provider
class AnthropicOAuthProvider(OAuthProvider):
    provider_name = "anthropic_oauth"
    model_prefixes = ["claude-code"]
    system_prompt_prefix = "You are Claude Code, Anthropic's official CLI for Claude."

    def __init__(self, credentials: dict[str, Any]):
        super().__init__(credentials)
        token = credentials.get("token")
        if not isinstance(token, str) or not token.strip():
            token = credentials.get("access")
        if not isinstance(token, str) or not token.strip():
            raise OAuthProviderError("Missing Anthropic OAuth access token in credentials")
        self._access_token = token
        self._client = anthropic.AsyncAnthropic(
            auth_token=self._access_token,
            max_retries=_MAX_RETRIES,
            timeout=_TIMEOUT_SECONDS,
        )

    async def stream_completion(
        self,
        model: str,
        messages: list[dict[str, Any]],
        *,
        tools: list[dict[str, Any]] | None = None,
        tool_choice: str | dict[str, Any] | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
        system_prompt: str | None = None,
        session_id: str | None = None,
        reasoning_effort: str | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[ProviderChunk]:
        _ = tool_choice, session_id, kwargs

        stream_kwargs: dict[str, Any] = {
            "model": self._extract_model_id(model),
            "messages": self._convert_messages(messages),
            "system": self._build_system_blocks(system_prompt),
            "max_tokens": int(max_tokens) if max_tokens is not None else _DEFAULT_MAX_TOKENS,
            "extra_headers": {
                "anthropic-beta": ",".join(_BETA_FEATURES),
                "anthropic-dangerous-direct-browser-access": "true",
                "user-agent": "claude-cli/2.1.2 (external, cli)",
                "x-app": "cli",
            },
        }
        if temperature is not None:
            stream_kwargs["temperature"] = temperature

        normalized_effort = _normalize_reasoning_effort(reasoning_effort)
        if normalized_effort != "none":
            max_tokens_value = int(stream_kwargs["max_tokens"])
            target_budget = _THINKING_BUDGETS.get(normalized_effort, _THINKING_BUDGETS["medium"])
            max_budget = max(256, max_tokens_value - 256)
            stream_kwargs["thinking"] = {
                "type": "enabled",
                "budget_tokens": min(target_budget, max_budget),
            }

        if tools:
            converted_tools = self._convert_tools(tools)
            if converted_tools:
                stream_kwargs["tools"] = converted_tools

        try:
            self._current_tool_index = 0
            self._tool_use_count = 0
            async with self._client.messages.stream(**stream_kwargs) as stream:
                async for event in stream:
                    chunk = self._map_event(event)
                    if chunk is not None:
                        yield chunk

                final_message = await stream.get_final_message()
                usage = self._extract_usage(final_message)
                stop_reason = final_message.stop_reason or "stop"
                yield self._final_chunk(stop_reason, usage)
        except anthropic.APIError as exc:
            raise OAuthProviderError(str(exc)) from exc

    def _map_event(self, event: Any) -> ProviderChunk | None:
        event_type = event.type

        if event_type == "text":
            if event.text:
                return self._text_chunk(event.text)
            return None

        if event_type == "content_block_start":
            content_block = event.content_block
            if getattr(content_block, "type", None) == "tool_use":
                # Use sequential counter rather than content_block index so that
                # thinking blocks (which precede tool_use and are never emitted as
                # chunks) don't create phantom empty slots in _consume_stream.
                index = self._tool_use_count
                self._tool_use_count += 1
                self._current_tool_index = index
                return self._tool_call_chunk(
                    index,
                    getattr(content_block, "id", None),
                    getattr(content_block, "name", None),
                    None,
                )
            return None

        if event_type == "input_json":
            return self._tool_call_chunk(self._current_tool_index, None, None, event.partial_json)

        if event_type == "error":
            error_msg = str(event) if event else "Anthropic OAuth stream error"
            raise OAuthProviderError(error_msg)

        # thinking, signature, content_block_stop, message_start, message_delta,
        # message_stop — all ignored; final chunk emitted after stream ends
        return None

    @staticmethod
    def _extract_usage(message: Any) -> dict[str, int] | None:
        usage = getattr(message, "usage", None)
        if usage is None:
            return None
        prompt_tokens = getattr(usage, "input_tokens", 0) or 0
        completion_tokens = getattr(usage, "output_tokens", 0) or 0
        return {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens,
        }

    @staticmethod
    def _extract_model_id(model: str) -> str:
        if model.lower().startswith("claude-code/"):
            return model.split("/", 1)[1]
        return model

    def _build_system_blocks(self, system_prompt: str | None) -> list[dict[str, Any]]:
        blocks: list[dict[str, Any]] = [
            {
                "type": "text",
                "text": self.system_prompt_prefix,
                "cache_control": {"type": "ephemeral"},
            }
        ]
        if system_prompt:
            blocks.append(
                {
                    "type": "text",
                    "text": _sanitize_text(system_prompt),
                    "cache_control": {"type": "ephemeral"},
                }
            )
        return blocks

    def _convert_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
        converted: list[dict[str, Any]] = []

        for message in messages:
            role = str(message.get("role", "user"))
            content = message.get("content", "")

            if role == "system":
                continue

            if role == "tool":
                converted.append(
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "tool_result",
                                "tool_use_id": str(message.get("tool_call_id", "")),
                                "content": _sanitize_text(str(content)),
                            }
                        ],
                    }
                )
                continue

            if role == "assistant":
                assistant_blocks: list[dict[str, Any]] = []
                assistant_blocks.extend(self._convert_content_blocks(content))

                tool_calls = message.get("tool_calls")
                if isinstance(tool_calls, list):
                    for tool_call in tool_calls:
                        if not isinstance(tool_call, dict):
                            continue
                        function = tool_call.get("function")
                        function_data = function if isinstance(function, dict) else {}
                        raw_arguments = str(function_data.get("arguments", "{}"))
                        try:
                            tool_input = json.loads(raw_arguments)
                        except json.JSONDecodeError:
                            tool_input = {}
                        if not isinstance(tool_input, dict):
                            tool_input = {}

                        assistant_blocks.append(
                            {
                                "type": "tool_use",
                                "id": str(tool_call.get("id", "")),
                                "name": str(function_data.get("name", "")),
                                "input": tool_input,
                            }
                        )

                if assistant_blocks:
                    converted.append({"role": "assistant", "content": assistant_blocks})
                continue

            target_role = "assistant" if role == "assistant" else "user"
            if isinstance(content, str):
                if content:
                    converted.append({"role": target_role, "content": _sanitize_text(content)})
                continue

            blocks = self._convert_content_blocks(content)
            if blocks:
                converted.append({"role": target_role, "content": blocks})

        return converted

    @staticmethod
    def _convert_content_blocks(content: Any) -> list[dict[str, Any]]:
        blocks: list[dict[str, Any]] = []
        if isinstance(content, str):
            if content:
                blocks.append({"type": "text", "text": _sanitize_text(content)})
            return blocks

        if not isinstance(content, list):
            return blocks

        for item in content:
            if not isinstance(item, dict):
                continue
            if item.get("type") == "text":
                text = str(item.get("text", ""))
                if text:
                    blocks.append({"type": "text", "text": _sanitize_text(text)})
            elif item.get("type") == "image_url":
                image_obj = item.get("image_url")
                if not isinstance(image_obj, dict):
                    continue
                url = str(image_obj.get("url", ""))
                if not url.startswith("data:"):
                    continue
                metadata, _, encoded = url.partition(",")
                media_type = metadata.split(";")[0].replace("data:", "") or "image/png"
                if encoded:
                    blocks.append(
                        {
                            "type": "image",
                            "source": {
                                "type": "base64",
                                "media_type": media_type,
                                "data": encoded,
                            },
                        }
                    )
        return blocks

    @staticmethod
    def _convert_tools(tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
        converted: list[dict[str, Any]] = []
        for tool in tools:
            if tool.get("type") != "function":
                continue
            function = tool.get("function")
            function_data = function if isinstance(function, dict) else {}
            converted.append(
                {
                    "name": str(function_data.get("name", "")),
                    "description": str(function_data.get("description", "")),
                    "input_schema": function_data.get(
                        "parameters",
                        {"type": "object", "properties": {}},
                    ),
                }
            )
        return converted
