"""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.credentials import Credential
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
    _credential: Credential
    _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.
        """

    async def aclose(self) -> None:
        """Close the underlying async client, releasing connection resources.

        Subclasses that lazily create an async HTTP client should override
        this method if the default ``await self._client.close()`` is not
        appropriate.
        """
        if hasattr(self, '_client') and self._client is not None:
            if hasattr(self._client, 'close'):
                await self._client.close()
            self._client = None


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.

    Normalises generation parameters from a prompt config dict,
    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', ['top_p']),
        ('max_tokens', ['max_tokens']),
        ('stop', ['stop']),
        ('frequency_penalty', ['frequency_penalty']),
        ('presence_penalty', ['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
