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

import abc
import threading
from collections.abc import Callable
from typing import Any, Literal, TypeVar, cast, get_args

from pydantic import BaseModel, Field

from dotpromptz.credentials import Credential
from dotpromptz.typing import RenderedPrompt

ThinkingLevel = Literal['low', 'medium', 'high', 'adaptive']
DEFAULT_THINKING_LEVEL: ThinkingLevel = 'high'
ConvertedT = TypeVar('ConvertedT')


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 GeneratedImage(BaseModel):
    """Unified image artifact produced by provider image generation APIs.

    Attributes:
        data: Raw image bytes.
        mime_type: MIME type for ``data`` (for example ``image/png``).
    """

    data: bytes
    mime_type: str


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

    Attributes:
        text: The text content of the response.
        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.
        images: Generated image artifacts when ``output.format`` is ``image``.
    """

    text: str | None = None
    usage: Usage | None = None
    raw: Any = None
    model: str | None = None
    finish_reason: str | None = None
    images: list[GeneratedImage] | 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) -> RenderedPrompt:
        """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})

    def _convert_with_model(
        self,
        rendered: RenderedPrompt,
        converter: Callable[[RenderedPrompt], ConvertedT],
    ) -> ConvertedT:
        """Convert a rendered prompt after ensuring a model is present.

        Args:
            rendered: The rendered prompt to convert.
            converter: Provider-specific converter callable.

        Returns:
            Converter output.
        """
        return converter(self._ensure_model(rendered))

    def _get_or_create_client(self, factory: Callable[[], Any]) -> Any:
        """Return a lazily created SDK client in a thread-safe manner.

        Args:
            factory: Callable used to create the underlying SDK client.

        Returns:
            The cached SDK client instance.
        """
        if self._client is None:
            with self._client_lock:
                if self._client is None:
                    self._client = factory()
        return self._client

    @abc.abstractmethod
    async def generate(
        self,
        rendered: RenderedPrompt,
        **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 and usage info.
        """

    @abc.abstractmethod
    def convert(self, rendered: RenderedPrompt) -> 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_thinking_level(value: Any) -> ThinkingLevel:
    """Resolve a supported thinking level from an arbitrary config value.

    Args:
        value: Raw thinking value from config.

    Returns:
        A valid thinking level. Invalid or missing values fall back to
        ``DEFAULT_THINKING_LEVEL``.
    """
    if isinstance(value, str):
        level = value.strip().lower()
        if level in get_args(ThinkingLevel):
            return cast(ThinkingLevel, level)
    return DEFAULT_THINKING_LEVEL


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. ``thinking`` is always
        present and defaults to ``high``.
    """
    mapping: list[tuple[str, list[str]]] = [
        ('thinking', ['thinking']),
        ('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
    result['thinking'] = _resolve_thinking_level(result.get('thinking'))
    return result


def resolve_rendered_config_and_model(rendered: RenderedPrompt) -> tuple[dict[str, Any], str]:
    """Return a normalized config dict and required model from a prompt.

    Args:
        rendered: The rendered prompt to inspect.

    Returns:
        A tuple of ``(config, model)``.

    Raises:
        ValueError: If no model is specified.
    """
    config = 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.')
    return config, cast(str, resolved_model)
