"""xAI (Grok) provider using OpenAI-compatible API."""

from __future__ import annotations

from collections.abc import Generator
from typing import Any, Final

import httpx
import openai
from tenacity import retry, stop_after_attempt, wait_exponential

from ember._internal.exceptions import ProviderAPIError
from ember.models.providers.base import BaseProvider
from ember.models.schemas import ChatResponse, UsageStats

_XAI_BASE_URL: Final[str] = "https://api.x.ai/v1"

_SUPPORTED_MODELS: Final[set[str]] = {
    "grok-3",
    "grok-3-fast",
    "grok-4",
    "grok-4-fast",
    "grok-4.1-fast",
    "grok-code-fast-1",
}


def _resolve_xai_api_key() -> str:
    """Resolve xAI API key from credential store."""
    from ember._internal.context.runtime import EmberContext
    from ember.core.credentials import CredentialNotFoundError

    try:
        return EmberContext.current().get_credential("xai")
    except CredentialNotFoundError as exc:
        raise ValueError(
            "xAI credential not configured. "
            "Run: ember configure set providers.xai.api_key YOUR_KEY"
        ) from exc


def _create_xai_client(api_key: str) -> openai.OpenAI:
    """Create OpenAI client configured for xAI API."""
    from ember._internal.context.runtime import EmberContext

    if not api_key or not api_key.strip():
        raise ValueError("xAI api_key must be a non-empty string")

    base_url = EmberContext.current().get_config("providers.xai.base_url")
    if not isinstance(base_url, str) or not base_url.strip():
        base_url = _XAI_BASE_URL

    return openai.OpenAI(
        api_key=api_key,
        base_url=base_url,
        max_retries=0,
        timeout=httpx.Timeout(120.0, connect=5.0),
    )


class XAIProvider(BaseProvider):
    """xAI provider for Grok models via OpenAI-compatible API."""

    supports_responses_api = False

    def __init__(self, api_key: str | None = None):
        super().__init__(api_key)
        self._client = _create_xai_client(self.api_key)

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=30),
        reraise=True,
    )
    def complete(self, prompt: str, model: str, **kwargs: Any) -> ChatResponse:
        if kwargs.pop("stream", False):
            raise ProviderAPIError(
                "Streaming not implemented; use stream_complete()",
                context={"model": model},
            )

        system_prompt = kwargs.pop("system", None)
        temperature = kwargs.pop("temperature", None)
        max_tokens = kwargs.pop("max_tokens", None)

        messages: list[dict[str, str]] = []
        if system_prompt:
            messages.append({"role": "system", "content": str(system_prompt)})
        messages.append({"role": "user", "content": prompt})

        request_kwargs: dict[str, Any] = {"model": model, "messages": messages}
        if temperature is not None:
            request_kwargs["temperature"] = temperature
        if max_tokens is not None:
            request_kwargs["max_tokens"] = max_tokens

        try:
            response = self._client.chat.completions.create(**request_kwargs)
        except openai.APIError as exc:
            raise ProviderAPIError(
                f"xAI API error: {exc}",
                context={"model": model},
            ) from exc

        choice = response.choices[0]
        text = choice.message.content or ""

        usage = response.usage
        usage_stats = UsageStats(
            prompt_tokens=usage.prompt_tokens if usage else 0,
            completion_tokens=usage.completion_tokens if usage else 0,
            total_tokens=usage.total_tokens if usage else 0,
        )

        return ChatResponse(
            data=text,
            raw_output={"model": model, "finish_reason": choice.finish_reason},
            usage=usage_stats,
            model_id=response.model or model,
        )

    def stream_complete(
        self, prompt: str, model: str, **kwargs: Any
    ) -> Generator[str, None, ChatResponse]:
        system_prompt = kwargs.pop("system", None)
        temperature = kwargs.pop("temperature", None)
        max_tokens = kwargs.pop("max_tokens", None)

        messages: list[dict[str, str]] = []
        if system_prompt:
            messages.append({"role": "system", "content": str(system_prompt)})
        messages.append({"role": "user", "content": prompt})

        request_kwargs: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "stream": True,
        }
        if temperature is not None:
            request_kwargs["temperature"] = temperature
        if max_tokens is not None:
            request_kwargs["max_tokens"] = max_tokens

        try:
            stream = self._client.chat.completions.create(**request_kwargs)
        except openai.APIError as exc:
            raise ProviderAPIError(
                f"xAI API error: {exc}",
                context={"model": model},
            ) from exc

        text_chunks: list[str] = []
        resolved_model = model
        finish_reason: str | None = None
        usage_stats = UsageStats()

        for chunk in stream:
            if not chunk.choices:
                continue
            choice = chunk.choices[0]
            if choice.finish_reason:
                finish_reason = choice.finish_reason
            if chunk.model:
                resolved_model = chunk.model
            delta = choice.delta
            if delta.content:
                text_chunks.append(delta.content)
                yield delta.content

        # Capture usage from the final chunk if available
        if hasattr(stream, "_info") and stream._info:
            final_usage = getattr(stream._info, "usage", None)
            if final_usage:
                pt = final_usage.prompt_tokens
                ct = final_usage.completion_tokens
                tt = final_usage.total_tokens
                usage_stats = UsageStats(
                    prompt_tokens=pt if pt is not None else 0,
                    completion_tokens=ct if ct is not None else 0,
                    total_tokens=tt if tt is not None else 0,
                )

        return ChatResponse(
            data="".join(text_chunks),
            raw_output={"model": resolved_model, "finish_reason": finish_reason},
            usage=usage_stats,
            model_id=resolved_model,
        )
