"""Anthropic Claude adapter: conversion helpers and async API wrapper."""

import base64
import json
import threading
from typing import Any

from anthropic import AsyncAnthropic

from dotpromptz.image import resolve_image_ref
from dotpromptz.typing import (
    DataPart,
    MediaPart,
    Message,
    Part,
    RenderedPrompt,
    Role,
    TextPart,
)

from ._base import (
    Adapter,
    Credential,
    GenerateResponse,
    Usage,
    resolve_config_params,
    resolve_rendered_config_and_model,
)

# ---------------------------------------------------------------------------
# Conversion helpers
# ---------------------------------------------------------------------------

_ANTHROPIC_MIN_THINKING_BUDGET = 1024
_ANTHROPIC_THINKING_REQUIRED_TEMPERATURE = 1
_ANTHROPIC_DEFAULT_MAX_TOKENS = 16_000
_ANTHROPIC_ADAPTIVE_THINKING_MODEL = 'claude-opus-4-6'
_ANTHROPIC_THINKING_BUDGET_RATIOS = {
    'low': 0.25,
    'medium': 0.5,
    'high': 1.0,
}


def _to_anthropic_thinking_config(model: str, level: str, max_tokens: int) -> dict[str, Any]:
    """Map a canonical thinking level to Anthropic thinking config.

    Args:
        model: Anthropic model ID for this request.
        level: Canonical thinking level (low/medium/high).
        max_tokens: Effective Anthropic ``max_tokens`` for this request.

    Returns:
        Anthropic ``thinking`` config.

    Raises:
        ValueError: If ``max_tokens`` is too small for Anthropic thinking.
    """
    if model == _ANTHROPIC_ADAPTIVE_THINKING_MODEL:
        # Anthropic recommends adaptive thinking for this model.
        return {'type': 'adaptive'}

    max_budget = max_tokens - 1
    if max_budget < _ANTHROPIC_MIN_THINKING_BUDGET:
        raise ValueError(
            f'Anthropic thinking requires max_tokens > {_ANTHROPIC_MIN_THINKING_BUDGET}; got {max_tokens}.'
        )
    budget = int(max_budget * _ANTHROPIC_THINKING_BUDGET_RATIOS.get(level, 1.0))

    budget = max(_ANTHROPIC_MIN_THINKING_BUDGET, min(budget, max_budget))
    return {'type': 'enabled', 'budget_tokens': budget}


def _part_to_anthropic_block(part: Part) -> dict[str, Any] | None:
    """Convert a dotprompt ``Part`` to an Anthropic content block.

    Args:
        part: A dotprompt content part.

    Returns:
        A dict representing an Anthropic content block, or ``None``.
    """
    if isinstance(part, TextPart):
        return {'type': 'text', 'text': part.text}
    if isinstance(part, MediaPart):
        resolved = resolve_image_ref(
            part.media.url,
            content_type=part.media.content_type,
            read_bytes=True,
            read_remote_bytes=True,
        )
        if resolved.data is None:
            raise ValueError('Image bytes are missing during Anthropic request conversion.')
        encoded = base64.b64encode(resolved.data).decode('ascii')
        return {
            'type': 'image',
            'source': {
                'type': 'base64',
                'media_type': resolved.mime_type,
                'data': encoded,
            },
        }
    if isinstance(part, DataPart):
        return {'type': 'text', 'text': json.dumps(part.data, ensure_ascii=False)}
    return None


def _message_to_anthropic(msg: Message) -> dict[str, Any]:
    """Convert a dotprompt ``Message`` to a raw Anthropic message dict.

    Note: system messages should be extracted separately — Anthropic puts
    them in the top-level ``system`` parameter, not in ``messages``.

    Args:
        msg: A dotprompt message.

    Returns:
        A dict with ``role`` and ``content`` suitable for the Anthropic API.
    """
    role = 'assistant' if msg.role == Role.MODEL else msg.role.value

    blocks: list[dict[str, Any]] = []
    for part in msg.content:
        block = _part_to_anthropic_block(part)
        if block is not None:
            blocks.append(block)

    # Simplify single text block to a string.
    content: str | list[dict[str, Any]] = blocks
    if len(blocks) == 1 and blocks[0].get('type') == 'text':
        content = blocks[0]['text']

    return {'role': role, 'content': content}


def to_anthropic_messages(messages: list[Message]) -> tuple[str | None, list[dict[str, Any]]]:
    """Convert dotprompt messages to Anthropic format.

    Extracts system messages separately (Anthropic puts them at top level).

    Args:
        messages: List of dotprompt ``Message`` objects.

    Returns:
        A tuple of ``(system_text, message_list)``.
    """
    system_parts: list[str] = []
    api_messages: list[dict[str, Any]] = []

    for msg in messages:
        if msg.role == Role.SYSTEM:
            for part in msg.content:
                if isinstance(part, TextPart):
                    system_parts.append(part.text)
        else:
            api_messages.append(_message_to_anthropic(msg))

    system_text = '\n'.join(system_parts) if system_parts else None
    return system_text, api_messages


def to_anthropic_request(
    rendered: RenderedPrompt,
    *,
    default_max_tokens: int = _ANTHROPIC_DEFAULT_MAX_TOKENS,
) -> dict[str, Any]:
    """Convert a ``RenderedPrompt`` to an Anthropic request dict.

    Args:
        rendered: The rendered prompt.
        default_max_tokens: Fallback ``max_tokens`` when prompt config omits it.

    Returns:
        A dict suitable for ``client.messages.create(**d)``.

    Raises:
        ValueError: If no model is specified.
    """
    system_text, messages = to_anthropic_messages(rendered.messages)
    config, resolved_model = resolve_rendered_config_and_model(rendered)

    request: dict[str, Any] = {
        'model': resolved_model,
        'messages': messages,
    }

    # Resolve generation parameters via centralised helper.
    params = resolve_config_params(config)

    # max_tokens is required by Anthropic — use resolved value or default.
    fallback_max_tokens = default_max_tokens if default_max_tokens > 0 else _ANTHROPIC_DEFAULT_MAX_TOKENS
    raw_max_tokens = params.get('max_tokens', fallback_max_tokens)
    max_tokens = raw_max_tokens if isinstance(raw_max_tokens, int) and raw_max_tokens > 0 else fallback_max_tokens
    request['max_tokens'] = max_tokens

    request['thinking'] = _to_anthropic_thinking_config(resolved_model, params.get('thinking', 'high'), max_tokens)

    if system_text:
        request['system'] = system_text

    if params.get('temperature') is not None:
        temperature = params['temperature']
        if request.get('thinking') and temperature != _ANTHROPIC_THINKING_REQUIRED_TEMPERATURE:
            request['temperature'] = _ANTHROPIC_THINKING_REQUIRED_TEMPERATURE
        else:
            request['temperature'] = temperature
    if params.get('top_p') is not None:
        request['top_p'] = params['top_p']
    stop = params.get('stop')
    if stop:
        request['stop_sequences'] = stop if isinstance(stop, list) else [stop]

    if rendered.output and rendered.output.format in ('json', 'yaml') and rendered.output.schema_value is not None:
        request['output_config'] = {
            'format': {
                'type': 'json_schema',
                'schema': rendered.output.schema_value,
            },
        }

    return request


# ---------------------------------------------------------------------------
# Anthropic Adapter
# ---------------------------------------------------------------------------


class AnthropicAdapter(Adapter):
    """Adapter for the Anthropic Messages API (Claude).

    Args:
        credential: A ``Credential`` with ``api_key`` and optional ``base_url``.
        default_model: Fallback model name.
        max_tokens: Default max tokens when not in prompt config.
    """

    def __init__(
        self,
        *,
        credential: Credential,
        default_model: str = 'claude-sonnet-4-20250514',
        max_tokens: int = _ANTHROPIC_DEFAULT_MAX_TOKENS,
    ) -> None:
        """Initialise the Anthropic adapter."""
        self._credential = credential
        self._default_model = default_model
        self._max_tokens = max_tokens
        self._client: Any = None
        self._client_lock = threading.Lock()

    def _get_client(self) -> Any:
        """Return (and lazily create) the ``AsyncAnthropic`` client."""

        def _factory() -> AsyncAnthropic:
            kwargs: dict[str, Any] = {'api_key': self._credential.api_key}
            if self._credential.base_url:
                kwargs['base_url'] = self._credential.base_url
            return AsyncAnthropic(**kwargs)

        return self._get_or_create_client(_factory)

    def convert(self, rendered: RenderedPrompt) -> dict[str, Any]:
        """Convert a ``RenderedPrompt`` to an Anthropic request dict.

        Args:
            rendered: The rendered prompt.

        Returns:
            A dict suitable for ``client.messages.create(**d)``.
        """

        def _to_request(prepared: RenderedPrompt) -> dict[str, Any]:
            return to_anthropic_request(prepared, default_max_tokens=self._max_tokens)

        return self._convert_with_model(rendered, _to_request)

    async def generate(
        self,
        rendered: RenderedPrompt,
        **kwargs: Any,
    ) -> GenerateResponse:
        """Call the Anthropic Messages API.

        Args:
            rendered: The rendered prompt.
            **kwargs: Extra keyword args forwarded to the SDK.

        Returns:
            A ``GenerateResponse``.
        """
        client = self._get_client()
        request_dict = self.convert(rendered)
        request_dict.update(kwargs)

        response = await client.messages.create(**request_dict)

        text_parts: list[str] = []
        for block in response.content:
            if block.type == 'text':
                text_parts.append(block.text)

        usage: Usage | None = None
        if response.usage:
            usage = Usage(
                prompt_tokens=response.usage.input_tokens,
                completion_tokens=response.usage.output_tokens,
                total_tokens=response.usage.input_tokens + response.usage.output_tokens,
            )

        return GenerateResponse(
            text='\n'.join(text_parts) if text_parts else None,
            usage=usage,
            raw=response,
            model=response.model,
            finish_reason=response.stop_reason,
        )
