"""Provider-agnostic entry point for working with large language models."""

from __future__ import annotations

import asyncio
from typing import Dict, Optional, Protocol, Sequence, Union

from .types import EmbeddingVector, ImageInput, LLMResponse, MessageSequence
from .utils import run_sync_in_thread


class SupportsGenerateResponse(Protocol):
    """Protocol describing provider clients."""

    def generate_response(
        self,
        *,
        api_key: str,
        prompt: Optional[str] = None,
        model: str,
        max_tokens: int = 32000,
        reasoning_effort: Optional[str] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        images: Optional[Sequence[ImageInput]] = None,
        messages: Optional[MessageSequence] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> str:
        ...

    async def async_generate_response(
        self,
        *,
        api_key: str,
        prompt: Optional[str] = None,
        model: str,
        max_tokens: int = 32000,
        reasoning_effort: Optional[str] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        images: Optional[Sequence[ImageInput]] = None,
        messages: Optional[MessageSequence] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> str:
        ...

    def generate_image(
        self,
        *,
        api_key: str,
        prompt: str,
        model: str,
        image_size: Optional[str] = None,
        aspect_ratio: Optional[str] = None,
        image: Optional[ImageInput] = None,
    ) -> bytes:
        ...

    async def async_generate_image(
        self,
        *,
        api_key: str,
        prompt: str,
        model: str,
        image_size: Optional[str] = None,
        aspect_ratio: Optional[str] = None,
        image: Optional[ImageInput] = None,
    ) -> bytes:
        ...

    def list_models(
        self,
        *,
        api_key: str,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> Sequence[dict[str, Optional[str]]]:
        ...

    async def async_list_models(
        self,
        *,
        api_key: str,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> Sequence[dict[str, Optional[str]]]:
        ...


class SupportsEmbeddings(Protocol):
    def embed_content(
        self,
        *,
        api_key: str,
        model: str,
        contents: Union[str, Sequence[str]],
        task_type: Optional[str] = None,
        output_dimensionality: Optional[int] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> list[EmbeddingVector]:
        ...

    async def async_embed_content(
        self,
        *,
        api_key: str,
        model: str,
        contents: Union[str, Sequence[str]],
        task_type: Optional[str] = None,
        output_dimensionality: Optional[int] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> list[EmbeddingVector]:
        ...


class LLMClient:
    """Central client capable of routing requests to different providers."""

    def __init__(
        self,
        providers: Optional[Dict[str, SupportsGenerateResponse]] = None,
    ) -> None:
        self._providers: Dict[str, SupportsGenerateResponse] = {}
        default_providers = providers or self._discover_default_providers()
        for name, client in default_providers.items():
            self.register_provider(name, client)

        if not self._providers:
            raise RuntimeError(
                "No provider implementations registered. Install the required extras "
                "for your target provider (e.g. `pip install openai`)."
            )

    def register_provider(self, name: str, client: SupportsGenerateResponse) -> None:
        """Register or overwrite a provider implementation."""
        if not name:
            raise ValueError("Provider name must be provided.")
        if client is None:
            raise ValueError("Provider client must be provided.")

        self._providers[name.lower()] = client

    def generate_response(
        self,
        *,
        provider: str,
        api_key: str,
        prompt: Optional[str] = None,
        model: str,
        max_tokens: int = 32000,
        reasoning_effort: Optional[str] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        images: Optional[Sequence[ImageInput]] = None,
        messages: Optional[MessageSequence] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> str:
        """Generate a response using the selected provider."""
        if not provider:
            raise ValueError("provider must be provided.")

        provider_client = self._providers.get(provider.lower())
        if provider_client is None:
            available = ", ".join(sorted(self._providers)) or "<none>"
            raise ValueError(
                f"Unknown provider '{provider}'. Available providers: {available}."
            )

        return provider_client.generate_response(
            api_key=api_key,
            prompt=prompt,
            model=model,
            max_tokens=max_tokens,
            reasoning_effort=reasoning_effort,
            temperature=temperature,
            top_p=top_p,
            images=images,
            messages=messages,
            request_id=request_id,
            timeout_s=timeout_s,
            max_retries=max_retries,
            retry_backoff_s=retry_backoff_s,
        )

    def generate_response_with_usage(
        self,
        *,
        provider: str,
        api_key: str,
        prompt: Optional[str] = None,
        model: str,
        max_tokens: int = 32000,
        reasoning_effort: Optional[str] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        images: Optional[Sequence[ImageInput]] = None,
        messages: Optional[MessageSequence] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> LLMResponse:
        """Generate a response using the selected provider (with token usage when available)."""
        if not provider:
            raise ValueError("provider must be provided.")

        provider_client = self._providers.get(provider.lower())
        if provider_client is None:
            available = ", ".join(sorted(self._providers)) or "<none>"
            raise ValueError(
                f"Unknown provider '{provider}'. Available providers: {available}."
            )

        method = getattr(provider_client, "generate_response_with_usage", None)
        if callable(method):
            return method(
                api_key=api_key,
                prompt=prompt,
                model=model,
                max_tokens=max_tokens,
                reasoning_effort=reasoning_effort,
                temperature=temperature,
                top_p=top_p,
                images=images,
                messages=messages,
                request_id=request_id,
                timeout_s=timeout_s,
                max_retries=max_retries,
                retry_backoff_s=retry_backoff_s,
            )

        text = provider_client.generate_response(
            api_key=api_key,
            prompt=prompt,
            model=model,
            max_tokens=max_tokens,
            reasoning_effort=reasoning_effort,
            temperature=temperature,
            top_p=top_p,
            images=images,
            messages=messages,
            request_id=request_id,
            timeout_s=timeout_s,
            max_retries=max_retries,
            retry_backoff_s=retry_backoff_s,
        )
        return LLMResponse(text=text, usage=None, provider=provider, model=model)

    async def async_generate_response(
        self,
        *,
        provider: str,
        api_key: str,
        prompt: Optional[str] = None,
        model: str,
        max_tokens: int = 32000,
        reasoning_effort: Optional[str] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        images: Optional[Sequence[ImageInput]] = None,
        messages: Optional[MessageSequence] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> str:
        """Generate a response using the selected provider (async)."""
        if not provider:
            raise ValueError("provider must be provided.")

        provider_client = self._providers.get(provider.lower())
        if provider_client is None:
            available = ", ".join(sorted(self._providers)) or "<none>"
            raise ValueError(
                f"Unknown provider '{provider}'. Available providers: {available}."
            )

        async_method = getattr(provider_client, "async_generate_response", None)
        if async_method is not None and asyncio.iscoroutinefunction(async_method):
            return await async_method(
                api_key=api_key,
                prompt=prompt,
                model=model,
                max_tokens=max_tokens,
                reasoning_effort=reasoning_effort,
                temperature=temperature,
                top_p=top_p,
                images=images,
                messages=messages,
                request_id=request_id,
                timeout_s=timeout_s,
                max_retries=max_retries,
                retry_backoff_s=retry_backoff_s,
            )

        return await run_sync_in_thread(
            lambda: provider_client.generate_response(
                api_key=api_key,
                prompt=prompt,
                model=model,
                max_tokens=max_tokens,
                reasoning_effort=reasoning_effort,
                temperature=temperature,
                top_p=top_p,
                images=images,
                messages=messages,
                request_id=request_id,
                timeout_s=timeout_s,
                max_retries=max_retries,
                retry_backoff_s=retry_backoff_s,
            )
        )

    def generate_image(
        self,
        *,
        provider: str,
        api_key: str,
        prompt: str,
        model: str,
        image_size: Optional[str] = None,
        aspect_ratio: Optional[str] = None,
        image: Optional[ImageInput] = None,
    ) -> bytes:
        """Generate an image using the selected provider."""
        if not provider:
            raise ValueError("provider must be provided.")

        provider_client = self._providers.get(provider.lower())
        if provider_client is None:
            available = ", ".join(sorted(self._providers)) or "<none>"
            raise ValueError(
                f"Unknown provider '{provider}'. Available providers: {available}."
            )

        return provider_client.generate_image(
            api_key=api_key,
            prompt=prompt,
            model=model,
            image_size=image_size,
            aspect_ratio=aspect_ratio,
            image=image,
        )

    async def async_generate_image(
        self,
        *,
        provider: str,
        api_key: str,
        prompt: str,
        model: str,
        image_size: Optional[str] = None,
        aspect_ratio: Optional[str] = None,
        image: Optional[ImageInput] = None,
    ) -> bytes:
        """Generate an image using the selected provider (async)."""
        if not provider:
            raise ValueError("provider must be provided.")

        provider_client = self._providers.get(provider.lower())
        if provider_client is None:
            available = ", ".join(sorted(self._providers)) or "<none>"
            raise ValueError(
                f"Unknown provider '{provider}'. Available providers: {available}."
            )

        async_method = getattr(provider_client, "async_generate_image", None)
        if async_method is not None and asyncio.iscoroutinefunction(async_method):
            return await async_method(
                api_key=api_key,
                prompt=prompt,
                model=model,
                image_size=image_size,
                aspect_ratio=aspect_ratio,
                image=image,
            )

        return await run_sync_in_thread(
            lambda: provider_client.generate_image(
                api_key=api_key,
                prompt=prompt,
                model=model,
                image_size=image_size,
                aspect_ratio=aspect_ratio,
                image=image,
            )
        )

    def list_models(
        self,
        *,
        provider: str,
        api_key: str,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> Sequence[dict[str, Optional[str]]]:
        """List models available for the specified provider."""
        if not provider:
            raise ValueError("provider must be provided.")

        provider_client = self._providers.get(provider.lower())
        if provider_client is None:
            available = ", ".join(sorted(self._providers)) or "<none>"
            raise ValueError(
                f"Unknown provider '{provider}'. Available providers: {available}."
            )

        return provider_client.list_models(
            api_key=api_key,
            request_id=request_id,
            timeout_s=timeout_s,
            max_retries=max_retries,
            retry_backoff_s=retry_backoff_s,
        )

    async def async_list_models(
        self,
        *,
        provider: str,
        api_key: str,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> Sequence[dict[str, Optional[str]]]:
        """List models available for the specified provider (async)."""
        if not provider:
            raise ValueError("provider must be provided.")

        provider_client = self._providers.get(provider.lower())
        if provider_client is None:
            available = ", ".join(sorted(self._providers)) or "<none>"
            raise ValueError(
                f"Unknown provider '{provider}'. Available providers: {available}."
            )

        async_method = getattr(provider_client, "async_list_models", None)
        if async_method is not None and asyncio.iscoroutinefunction(async_method):
            return await async_method(
                api_key=api_key,
                request_id=request_id,
                timeout_s=timeout_s,
                max_retries=max_retries,
                retry_backoff_s=retry_backoff_s,
            )

        return await run_sync_in_thread(
            lambda: provider_client.list_models(
                api_key=api_key,
                request_id=request_id,
                timeout_s=timeout_s,
                max_retries=max_retries,
                retry_backoff_s=retry_backoff_s,
            )
        )

    async def async_generate_response_with_usage(
        self,
        *,
        provider: str,
        api_key: str,
        prompt: Optional[str] = None,
        model: str,
        max_tokens: int = 32000,
        reasoning_effort: Optional[str] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        images: Optional[Sequence[ImageInput]] = None,
        messages: Optional[MessageSequence] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> LLMResponse:
        """Generate a response using the selected provider (async, with token usage when available)."""
        if not provider:
            raise ValueError("provider must be provided.")

        provider_client = self._providers.get(provider.lower())
        if provider_client is None:
            available = ", ".join(sorted(self._providers)) or "<none>"
            raise ValueError(
                f"Unknown provider '{provider}'. Available providers: {available}."
            )

        async_method = getattr(provider_client, "async_generate_response_with_usage", None)
        if async_method is not None and asyncio.iscoroutinefunction(async_method):
            return await async_method(
                api_key=api_key,
                prompt=prompt,
                model=model,
                max_tokens=max_tokens,
                reasoning_effort=reasoning_effort,
                temperature=temperature,
                top_p=top_p,
                images=images,
                messages=messages,
                request_id=request_id,
                timeout_s=timeout_s,
                max_retries=max_retries,
                retry_backoff_s=retry_backoff_s,
            )

        method = getattr(provider_client, "generate_response_with_usage", None)
        if callable(method):
            return await run_sync_in_thread(
                lambda: method(
                    api_key=api_key,
                    prompt=prompt,
                    model=model,
                    max_tokens=max_tokens,
                    reasoning_effort=reasoning_effort,
                    temperature=temperature,
                    top_p=top_p,
                    images=images,
                    messages=messages,
                    request_id=request_id,
                    timeout_s=timeout_s,
                    max_retries=max_retries,
                    retry_backoff_s=retry_backoff_s,
                )
            )

        text = await self.async_generate_response(
            provider=provider,
            api_key=api_key,
            prompt=prompt,
            model=model,
            max_tokens=max_tokens,
            reasoning_effort=reasoning_effort,
            temperature=temperature,
            top_p=top_p,
            images=images,
            messages=messages,
            request_id=request_id,
            timeout_s=timeout_s,
            max_retries=max_retries,
            retry_backoff_s=retry_backoff_s,
        )
        return LLMResponse(text=text, usage=None, provider=provider, model=model)

    def embed_content(
        self,
        *,
        provider: str,
        api_key: str,
        model: str,
        contents: Union[str, Sequence[str]],
        task_type: Optional[str] = None,
        output_dimensionality: Optional[int] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> list[EmbeddingVector]:
        """Generate embeddings using the selected provider."""
        if not provider:
            raise ValueError("provider must be provided.")

        provider_client = self._providers.get(provider.lower())
        if provider_client is None:
            available = ", ".join(sorted(self._providers)) or "<none>"
            raise ValueError(
                f"Unknown provider '{provider}'. Available providers: {available}."
            )

        embed_method = getattr(provider_client, "embed_content", None)
        if not callable(embed_method):
            raise ValueError(f"Provider '{provider}' does not support embeddings.")

        return embed_method(
            api_key=api_key,
            model=model,
            contents=contents,
            task_type=task_type,
            output_dimensionality=output_dimensionality,
            request_id=request_id,
            timeout_s=timeout_s,
            max_retries=max_retries,
            retry_backoff_s=retry_backoff_s,
        )

    async def async_embed_content(
        self,
        *,
        provider: str,
        api_key: str,
        model: str,
        contents: Union[str, Sequence[str]],
        task_type: Optional[str] = None,
        output_dimensionality: Optional[int] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> list[EmbeddingVector]:
        """Generate embeddings using the selected provider (async)."""
        if not provider:
            raise ValueError("provider must be provided.")

        provider_client = self._providers.get(provider.lower())
        if provider_client is None:
            available = ", ".join(sorted(self._providers)) or "<none>"
            raise ValueError(
                f"Unknown provider '{provider}'. Available providers: {available}."
            )

        async_method = getattr(provider_client, "async_embed_content", None)
        if async_method is not None and asyncio.iscoroutinefunction(async_method):
            return await async_method(
                api_key=api_key,
                model=model,
                contents=contents,
                task_type=task_type,
                output_dimensionality=output_dimensionality,
                request_id=request_id,
                timeout_s=timeout_s,
                max_retries=max_retries,
                retry_backoff_s=retry_backoff_s,
            )

        embed_method = getattr(provider_client, "embed_content", None)
        if not callable(embed_method):
            raise ValueError(f"Provider '{provider}' does not support embeddings.")

        return await run_sync_in_thread(
            lambda: embed_method(
                api_key=api_key,
                model=model,
                contents=contents,
                task_type=task_type,
                output_dimensionality=output_dimensionality,
                request_id=request_id,
                timeout_s=timeout_s,
                max_retries=max_retries,
                retry_backoff_s=retry_backoff_s,
            )
        )

    @staticmethod
    def _discover_default_providers() -> Dict[str, SupportsGenerateResponse]:
        providers: Dict[str, SupportsGenerateResponse] = {}
        try:
            from .openai_client import OpenAIResponsesClient  # type: ignore
        except ModuleNotFoundError as exc:
            if exc.name != "openai":
                raise
            return providers

        providers["openai"] = OpenAIResponsesClient()

        try:
            from .gemini_client import GeminiClient  # type: ignore
        except ModuleNotFoundError as exc:
            if exc.name not in {"google", "google.genai"}:
                raise
        else:
            providers["gemini"] = GeminiClient()

        try:
            from .anthropic_client import AnthropicClient  # type: ignore
        except ModuleNotFoundError as exc:
            if exc.name != "anthropic":
                raise
        else:
            providers["anthropic"] = AnthropicClient()

        try:
            from .grok_client import GrokClient  # type: ignore
        except ModuleNotFoundError as exc:
            if exc.name != "xai_sdk":
                raise
        else:
            grok_client = GrokClient()
            providers["grok"] = grok_client
            providers.setdefault("xai", grok_client)

        return providers
