"""Base provider interface for ai-query."""

from __future__ import annotations

import asyncio
import base64
import json
from abc import ABC, abstractmethod
from typing import Any, AsyncIterator, TYPE_CHECKING

from ai_query.types import (
    GenerateTextResult,
    Message,
    ProviderOptions,
    StreamChunk,
    ToolSet,
    Usage,
    EmbedResult,
    EmbedManyResult,
    EmbeddingUsage,
)

if TYPE_CHECKING:
    from ai_query.transport import HTTPTransport


class BaseProvider(ABC):
    """Abstract base class for AI providers.

    To create a new provider adapter:
    1. Subclass BaseProvider
    2. Implement the `generate` method
    3. Implement the `stream` method for streaming support
    """

    # Provider identifier (e.g., "openai", "anthropic", "google")
    name: str

    def __init__(
        self,
        api_key: str | None = None,
        transport: HTTPTransport | None = None,
        **kwargs: Any,
    ):
        """Initialize the provider.

        Args:
            api_key: API key for the provider. If None, will try to read from
                     environment variable (provider-specific).
            transport: HTTP transport to use. If None, will auto-detect based on
                       environment (aiohttp for standard Python, fetch for Workers).
            **kwargs: Additional provider-specific configuration.
        """
        self.api_key = api_key
        self.config = kwargs
        self._transport = transport

    @property
    def transport(self) -> HTTPTransport:
        """Get or create the HTTP transport.

        Returns the transport passed to __init__, or auto-detects the appropriate
        transport for the current environment.
        """
        if self._transport is None:
            from ai_query.transport import get_default_transport

            self._transport = get_default_transport()
        return self._transport

    @abstractmethod
    async def generate(
        self,
        *,
        model: str,
        messages: list[Message],
        tools: ToolSet | None = None,
        provider_options: ProviderOptions | None = None,
        **kwargs: Any,
    ) -> GenerateTextResult:
        """Generate text using the provider's API.

        Args:
            model: The model identifier (without provider prefix).
            messages: List of messages in the conversation.
            tools: Optional dict of tool definitions for tool calling.
            provider_options: Provider-specific options.
            **kwargs: Additional parameters (max_tokens, temperature, etc.).

        Returns:
            GenerateTextResult containing the generated text and metadata.
        """
        pass

    @abstractmethod
    async def stream(
        self,
        *,
        model: str,
        messages: list[Message],
        tools: ToolSet | None = None,
        provider_options: ProviderOptions | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[StreamChunk]:
        """Stream text using the provider's API.

        Args:
            model: The model identifier (without provider prefix).
            messages: List of messages in the conversation.
            provider_options: Provider-specific options.
            **kwargs: Additional parameters (max_tokens, temperature, etc.).

        Yields:
            StreamChunk objects containing text and metadata.
            The final chunk will have is_final=True with usage and finish_reason.
        """
        pass
        # Make this an async generator
        if False:
            yield StreamChunk()

    def get_provider_options(
        self, provider_options: ProviderOptions | None
    ) -> dict[str, Any]:
        """Extract options specific to this provider.

        Args:
            provider_options: Full provider options dict.

        Returns:
            Options for this specific provider, or empty dict.
        """
        if provider_options is None:
            return {}
        return provider_options.get(self.name, {})

    async def _fetch_resource_as_base64(self, url: str) -> tuple[str, str]:
        """Fetch a resource from URL and return as base64 with media type.

        Args:
            url: The URL to fetch.

        Returns:
            Tuple of (base64_data, media_type).
        """
        data_bytes, content_type = await self.transport.get(url)
        # Extract just the mime type (remove charset etc)
        media_type = content_type.split(";")[0].strip()
        return base64.b64encode(data_bytes).decode(), media_type

    def _parse_sse_line(self, line: bytes | str) -> str | None:
        """Parse an SSE line and extract the data payload.

        Args:
            line: Raw SSE line (bytes or str).

        Returns:
            The data string if line starts with "data: ", None otherwise.
        """
        if isinstance(line, bytes):
            line = line.decode("utf-8")
        line = line.strip()
        if not line:
            return None
        if line.startswith("data: "):
            return line[6:]
        return None

    def _parse_sse_json(self, line: bytes | str) -> dict[str, Any] | None:
        """Parse an SSE line and decode the JSON payload.

        Args:
            line: Raw SSE line (bytes or str).

        Returns:
            Parsed JSON dict, or None if line is not valid SSE data or invalid JSON.
        """
        data = self._parse_sse_line(line)
        if data is None:
            return None
        if data == "[DONE]":
            return None
        try:
            return json.loads(data)
        except json.JSONDecodeError:
            return None

    def _accumulate_usage(self, total: Usage, delta: Usage) -> None:
        """Accumulate usage statistics in-place.

        Args:
            total: The running total Usage object (modified in-place).
            delta: The delta Usage object to add.
        """
        total.input_tokens += delta.input_tokens
        total.output_tokens += delta.output_tokens
        total.cached_tokens += delta.cached_tokens
        total.total_tokens += delta.total_tokens

    async def embed(
        self,
        *,
        model: str,
        value: str,
        provider_options: ProviderOptions | None = None,
        **kwargs: Any,
    ) -> EmbedResult:
        """Generate an embedding for a single value.

        Args:
            model: The embedding model identifier.
            value: The text to embed.
            provider_options: Provider-specific options.
            **kwargs: Additional parameters (dimensions, etc.).

        Returns:
            EmbedResult containing the embedding vector and metadata.

        Raises:
            NotImplementedError: If the provider doesn't support embeddings.
        """
        raise NotImplementedError(f"{self.name} provider does not support embeddings")

    async def embed_many(
        self,
        *,
        model: str,
        values: list[str],
        provider_options: ProviderOptions | None = None,
        **kwargs: Any,
    ) -> EmbedManyResult:
        """Generate embeddings for multiple values.

        Default implementation calls embed() for each value in parallel.
        Providers should override this for more efficient batch processing
        if their API supports it natively.

        Args:
            model: The embedding model identifier.
            values: List of texts to embed.
            provider_options: Provider-specific options.
            **kwargs: Additional parameters (dimensions, etc.).

        Returns:
            EmbedManyResult containing embedding vectors and metadata.
        """
        # Run all embed calls in parallel
        tasks = [
            self.embed(
                model=model,
                value=value,
                provider_options=provider_options,
                **kwargs,
            )
            for value in values
        ]
        results = await asyncio.gather(*tasks)

        # Aggregate results
        embeddings = [r.embedding for r in results]
        total_tokens = sum(r.usage.tokens for r in results)

        return EmbedManyResult(
            values=values,
            embeddings=embeddings,
            usage=EmbeddingUsage(tokens=total_tokens),
        )
