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

import json
import threading
from typing import Any

from anthropic import AsyncAnthropic

from dotpromptz.typing import (
    DataPart,
    MediaPart,
    Message,
    Part,
    RenderedPrompt,
    Role,
    TextPart,
    ToolDefinition,
    ToolRequestPart,
    ToolResponsePart,
)

from ._base import Adapter, Credential, GenerateResponse, ToolCallResult, Usage, parse_data_uri, resolve_config_params

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


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):
        url = part.media.url
        media_type = part.media.content_type or 'image/png'
        # Handle base64 data URIs.
        parsed = parse_data_uri(url)
        if parsed:
            inferred_type, data = parsed
            return {
                'type': 'image',
                'source': {
                    'type': 'base64',
                    'media_type': inferred_type or media_type,
                    'data': data,
                },
            }
        return {
            'type': 'image',
            'source': {'type': 'url', 'url': url},
        }
    if isinstance(part, DataPart):
        return {'type': 'text', 'text': json.dumps(part.data, ensure_ascii=False)}
    if isinstance(part, ToolRequestPart):
        req = part.tool_request
        return {
            'type': 'tool_use',
            'id': req.ref or req.name,
            'name': req.name,
            'input': req.input or {},
        }
    if isinstance(part, ToolResponsePart):
        resp = part.tool_response
        return {
            'type': 'tool_result',
            'tool_use_id': resp.ref or resp.name,
            'content': json.dumps(resp.output, ensure_ascii=False) if resp.output is not None else '',
        }
    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
    # Anthropic uses 'user' role for tool results too.
    if msg.role == Role.TOOL:
        role = 'user'

    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_tools(tool_defs: list[ToolDefinition] | None) -> list[dict[str, Any]] | None:
    """Convert dotprompt ``ToolDefinition`` list to Anthropic tool format.

    Args:
        tool_defs: List of dotprompt tool definitions.

    Returns:
        List of Anthropic tool dicts, or ``None``.
    """
    if not tool_defs:
        return None
    return [
        {
            'name': td.name,
            'description': td.description or '',
            'input_schema': td.input_schema,
        }
        for td in tool_defs
    ]


def to_anthropic_request(rendered: RenderedPrompt[Any]) -> 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)``.

    Raises:
        ValueError: If no model is specified.
    """
    system_text, messages = to_anthropic_messages(rendered.messages)
    config: dict[str, Any] = rendered.config if isinstance(rendered.config, dict) else {}

    resolved_model = config.get('model')
    if not resolved_model:
        raise ValueError('No model specified in rendered prompt.')

    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.
    request['max_tokens'] = params.get('max_tokens', 4096)

    if system_text:
        request['system'] = system_text

    tools = to_anthropic_tools(rendered.tool_defs)
    if tools:
        request['tools'] = tools

    if params.get('temperature') is not None:
        request['temperature'] = params['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 == 'json' and rendered.output.schema_ is not None:
        request['output_config'] = {
            'format': {
                'type': 'json_schema',
                'schema': rendered.output.schema_,
            },
        }

    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 = 4096,
    ) -> 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."""
        if self._client is None:
            with self._client_lock:
                if self._client is None:
                    kwargs: dict[str, Any] = {'api_key': self._credential.api_key}
                    if self._credential.base_url:
                        kwargs['base_url'] = self._credential.base_url
                    self._client = AsyncAnthropic(**kwargs)
        return self._client

    async def aclose(self) -> None:
        """Close the underlying ``AsyncAnthropic`` client."""
        if self._client is not None:
            await self._client.close()
            self._client = None

    def convert(self, rendered: RenderedPrompt[Any]) -> 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)``.
        """
        rendered = self._ensure_model(rendered)
        return to_anthropic_request(rendered)

    async def generate(
        self,
        rendered: RenderedPrompt[Any],
        **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] = []
        tool_calls: list[ToolCallResult] = []
        for block in response.content:
            if block.type == 'text':
                text_parts.append(block.text)
            elif block.type == 'tool_use':
                tool_calls.append(
                    ToolCallResult(
                        id=block.id,
                        name=block.name,
                        arguments=block.input if isinstance(block.input, dict) else {},
                    )
                )

        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,
            tool_calls=tool_calls,
            usage=usage,
            raw=response,
            model=response.model,
            finish_reason=response.stop_reason,
        )
