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

from __future__ import annotations

import abc
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'``).
        saved_path: File path where the image was saved.
    """

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


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

    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)
