from __future__ import annotations as _annotations

import asyncio
import os
from typing import TYPE_CHECKING, Any, overload

from pydantic_ai import ModelProfile
from pydantic_ai.exceptions import UserError
from pydantic_ai.profiles.grok import grok_model_profile, grok_realtime_model_profile
from pydantic_ai.providers import Provider

if TYPE_CHECKING:
    from pydantic_ai.realtime import RealtimeModelProfile

try:
    from xai_sdk import AsyncClient
except ImportError as _import_error:
    raise ImportError(
        'Please install the `xai-sdk` package to use the xAI provider, '
        'you can use the `xai` optional group — `pip install "pydantic-ai-slim[xai]"`'
    ) from _import_error


class _LazyAsyncClient:
    """Wrapper that creates a fresh AsyncClient per event loop.

    gRPC async channels bind to the event loop at creation time. If the client
    is created outside an async context (e.g. at module level) and later used
    inside asyncio.run(), the loop will differ, causing RuntimeError.
    This wrapper defers client creation and recreates it when the loop changes.
    See https://github.com/grpc/grpc/issues/32480.
    """

    def __init__(self, **kwargs: Any) -> None:
        self._kwargs = kwargs
        self._client: AsyncClient | None = None
        self._event_loop: asyncio.AbstractEventLoop | None = None

    def get_client(self) -> AsyncClient:
        running_loop: asyncio.AbstractEventLoop | None = None
        try:
            running_loop = asyncio.get_running_loop()
        except RuntimeError:
            pass

        if self._client is None or (running_loop is not None and running_loop is not self._event_loop):
            self._client = AsyncClient(**self._kwargs)
            self._event_loop = running_loop

        return self._client


class XaiProvider(Provider[AsyncClient]):
    """Provider for xAI API (native xAI SDK)."""

    @property
    def name(self) -> str:
        return 'xai'

    @property
    def base_url(self) -> str:
        # Canonical pricing/identity label, not the transport host: the xAI SDK is gRPC and the actual
        # channel target is set via `api_host`. This URL is used for usage/price lookup and telemetry only.
        return 'https://api.x.ai/v1'

    @property
    def client(self) -> AsyncClient:
        if self._lazy_client is not None:
            return self._lazy_client.get_client()
        return self._client

    @property
    def api_key(self) -> str | None:
        """The resolved API key, or `None` when the provider was built from a pre-configured `xai_client`.

        The gRPC `AsyncClient` doesn't expose its key, so this returns the one resolved from the `api_key`
        argument or `XAI_API_KEY`. Used by transports that authenticate outside the SDK, e.g.
        [`XaiRealtimeModel`][pydantic_ai.realtime.xai.XaiRealtimeModel]'s WebSocket `Authorization` header.
        """
        return self._api_key

    @property
    def api_host(self) -> str | None:
        """The custom `api_host` this provider was configured with, or `None`.

        Read by [`XaiRealtimeModel`][pydantic_ai.realtime.xai.XaiRealtimeModel] to reject a custom host
        it can't yet honor: the realtime WebSocket derives its URL from `base_url`, not the gRPC channel
        target that `api_host` sets.
        """
        return self._api_host

    @staticmethod
    def model_profile(model_name: str) -> ModelProfile | None:
        return grok_model_profile(model_name)

    @staticmethod
    def realtime_model_profile(model_name: str) -> RealtimeModelProfile:
        return grok_realtime_model_profile(model_name)

    @overload
    def __init__(
        self,
        *,
        api_key: str | None = None,
        api_host: str | None = None,
        timeout: float | None = None,
        metadata: tuple[tuple[str, str], ...] | None = None,
    ) -> None: ...

    @overload
    def __init__(self, *, xai_client: AsyncClient) -> None: ...

    def __init__(
        self,
        *,
        api_key: str | None = None,
        api_host: str | None = None,
        timeout: float | None = None,
        metadata: tuple[tuple[str, str], ...] | None = None,
        xai_client: AsyncClient | None = None,
    ) -> None:
        """Create a new xAI provider.

        Args:
            api_key: The API key to use for authentication, if not provided, the `XAI_API_KEY` environment variable
                will be used if available.
            api_host: The API host to use for the xAI SDK client.
            timeout: The client-level default timeout for the xAI SDK client, in seconds, applied to all requests
                made through it. The xAI SDK does not support per-request timeouts, so `ModelSettings.timeout` is
                not supported and has no effect.
            metadata: gRPC metadata to attach to every request the xAI SDK client makes, forwarded to
                [`xai_sdk.AsyncClient`][xai_sdk.AsyncClient]. This is client-scoped, not per-request, so it applies
                to every request made through the provider. The canonical use is xAI prompt-cache sticky routing via
                `metadata=(('x-grok-conv-id', '<conversation-id>'),)`; see the
                [xAI prompt-caching docs](https://docs.x.ai/developers/advanced-api-usage/prompt-caching/maximizing-cache-hits).
                Because it is client-scoped, a provider configured with conversation-specific metadata (e.g. a fixed
                `x-grok-conv-id`) must not be shared between unrelated conversations. Ignored when `xai_client` is
                passed.
            xai_client: An existing `xai_sdk.AsyncClient` to use. This takes precedence over `api_key`, `api_host`,
                `timeout`, and `metadata`.
        """
        self._lazy_client: _LazyAsyncClient | None = None
        # Retained so transports authenticating outside the gRPC SDK (e.g. the realtime WebSocket) can
        # read it back; `None` when a pre-configured `xai_client` was passed, since its key isn't exposed.
        self._api_key: str | None = None
        # Retained so the realtime WebSocket (which derives its host from `base_url`, not the gRPC
        # channel target) can detect a custom `api_host` it can't yet honor and fail loudly. Like
        # `_api_key`, left `None` when a pre-configured `xai_client` was passed: the client takes
        # precedence and the ignored `api_host` argument must not make the realtime model fail.
        self._api_host: str | None = None
        if xai_client is not None:
            self._client = xai_client
        else:
            api_key = api_key or os.getenv('XAI_API_KEY')
            if not api_key:
                raise UserError(
                    'Set the `XAI_API_KEY` environment variable or pass it via `XaiProvider(api_key=...)`'
                    ' to use the xAI provider.'
                )
            self._api_key = api_key
            self._api_host = api_host
            client_kwargs: dict[str, str | float | tuple[tuple[str, str], ...]] = {'api_key': api_key}
            if api_host is not None:
                client_kwargs['api_host'] = api_host
            if timeout is not None:
                client_kwargs['timeout'] = timeout
            if metadata is not None:
                client_kwargs['metadata'] = metadata
            self._lazy_client = _LazyAsyncClient(**client_kwargs)
            self._client = None  # type: ignore[assignment]
