"""Base adapter protocol and shared types for LLM API adapters."""

import abc
import threading
from typing import Any

from pydantic import BaseModel, Field

from dotpromptz.typing import RenderedPrompt


class Usage(BaseModel):
    """Token usage information from an LLM API call.

    Attributes:
        prompt_tokens: Number of tokens in the prompt.
        completion_tokens: Number of tokens in the completion.
        total_tokens: Total tokens used.
    """

    prompt_tokens: int = Field(default=0)
    completion_tokens: int = Field(default=0)
    total_tokens: int = Field(default=0)


class ToolCallResult(BaseModel):
    """A tool call returned by the LLM.

    Attributes:
        id: Unique identifier for the tool call.
        name: Name of the tool to invoke.
        arguments: Parsed arguments dict for the tool.
    """

    id: str
    name: str
    arguments: dict[str, Any]


class ImageResult(BaseModel):
    """Image data from a generative model response.

    Attributes:
        data: Raw image bytes.
        mime_type: MIME type of the image (e.g. ``'image/png'``).
    """

    data: bytes
    mime_type: str = 'image/png'


class GenerateResponse(BaseModel):
    """Unified response from any LLM adapter.

    Attributes:
        text: The text content of the response (None if tool calls only).
        image: Image data if the model generated an image.
        tool_calls: Tool calls requested by the model.
        usage: Token usage statistics.
        raw: The raw SDK response object for advanced use.
        model: The model that actually served the request.
        finish_reason: Why the model stopped generating.
    """

    text: str | None = None
    image: ImageResult | None = None
    tool_calls: list[ToolCallResult] = Field(default_factory=list)
    usage: Usage | None = None
    raw: Any = None
    model: str | None = None
    finish_reason: str | None = None


class Adapter(abc.ABC):
    """Abstract base class that all LLM adapters must implement.

    Subclasses must implement ``generate`` (async) and ``convert``.
    """

    _default_model: str
    _client: Any
    _client_lock: threading.Lock

    def _ensure_model(self, rendered: RenderedPrompt[Any]) -> RenderedPrompt[Any]:
        """Return *rendered* with ``config['model']`` guaranteed to be set.

        If ``config['model']`` does not already provide a model name, the
        adapter's ``_default_model`` is injected into a shallow copy of config.

        Args:
            rendered: The rendered prompt.

        Returns:
            The original *rendered* if a model is already present,
            otherwise a copy with ``config['model']`` populated.
        """
        config = rendered.config if isinstance(rendered.config, dict) else {}
        if config.get('model'):
            return rendered
        new_config = {**config, 'model': self._default_model}
        return rendered.model_copy(update={'config': new_config})

    @abc.abstractmethod
    async def generate(
        self,
        rendered: RenderedPrompt[Any],
        **kwargs: Any,
    ) -> GenerateResponse:
        """Send a rendered prompt to the LLM and return a unified response.

        Args:
            rendered: The rendered prompt from ``Dotprompt.render()``.
            **kwargs: Provider-specific overrides.

        Returns:
            A ``GenerateResponse`` with text, tool calls, and usage info.
        """

    @abc.abstractmethod
    def convert(self, rendered: RenderedPrompt[Any]) -> dict[str, Any]:
        """Convert a ``RenderedPrompt`` to the provider's native request dict.

        Useful for debugging / dry-run without making an API call.

        Args:
            rendered: The rendered prompt from ``Dotprompt.render()``.

        Returns:
            A dict representing the provider-specific request body.
        """


def parse_data_uri(url: str) -> tuple[str, str] | None:
    """Parse a ``data:`` URI into its MIME type and raw data.

    Args:
        url: A URL string that may be a ``data:`` URI.

    Returns:
        A ``(mime_type, data)`` tuple if *url* is a data URI, otherwise ``None``.
    """
    if not url.startswith('data:'):
        return None
    header, _, data = url.partition(',')
    mime_type = header.split(';')[0].replace('data:', '')
    return (mime_type, data)



def _first_defined(*values: Any) -> Any:
    """Return the first value that is not ``None``.

    Unlike chaining with ``or``, this correctly preserves falsy values
    such as ``0``, ``0.0``, and ``''``.

    Args:
        *values: Candidate values to check.

    Returns:
        The first non-``None`` value, or ``None`` if all are ``None``.
    """
    return next((v for v in values if v is not None), None)


def resolve_config_params(config: dict[str, Any]) -> dict[str, Any]:
    """Extract and normalise generation parameters from a prompt config dict.

    Handles both camelCase (``topP``) and snake_case (``top_p``) keys,
    returning a canonical dict with snake_case keys.  Falsy-but-valid
    values like ``0`` and ``0.0`` are preserved correctly.

    Args:
        config: The raw config dict from a rendered prompt.

    Returns:
        A dict with resolved generation parameters (only non-``None`` entries).
    """
    mapping: list[tuple[str, list[str]]] = [
        ('temperature', ['temperature']),
        ('top_p', ['topP', 'top_p']),
        ('max_tokens', ['maxTokens', 'max_tokens']),
        ('stop', ['stop', 'stopSequences']),
        ('frequency_penalty', ['frequencyPenalty', 'frequency_penalty']),
        ('presence_penalty', ['presencePenalty', 'presence_penalty']),
        ('seed', ['seed']),
    ]
    result: dict[str, Any] = {}
    for canonical, keys in mapping:
        value = _first_defined(*(config.get(k) for k in keys))
        if value is not None:
            result[canonical] = value
    return result
