"""Thin wrapper around the xAI Grok chat API via the xai-sdk package."""

from __future__ import annotations

import base64
import mimetypes
from pathlib import Path
import logging
from typing import Optional, Sequence

from xai_sdk import Client
from xai_sdk.chat import image as chat_image
from xai_sdk.chat import user

from .types import ImageInput, MessageSequence, normalize_messages
from .utils import clamp_retries, run_sync_in_thread, run_with_retries

logger = logging.getLogger(__name__)


class GrokClient:
    """Convenience wrapper around the xAI Grok chat API."""

    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:
        """Generate a response from the specified Grok model.

        Args:
            api_key: API key used to authenticate with xAI.
            prompt: Natural-language instruction or query for the model.
            model: Identifier of the Grok model to target (for example, ``"grok-3"``).
            max_tokens: Cap for tokens in the generated response, defaults to 32000.
            reasoning_effort: Optional hint for reasoning-focused models (``"low"`` or ``"high"``).
            temperature: Optional sampling temperature.
            top_p: Optional nucleus sampling value.
            images: Optional collection of image references (local paths, URLs, or data URLs).
            messages: Optional list of chat-style messages (role/content).
            request_id: Optional request identifier for tracing/logging.
            timeout_s: Optional request timeout in seconds (best-effort).
            max_retries: Optional retry count for transient failures.
            retry_backoff_s: Base delay (seconds) for exponential backoff between retries.

        Returns:
            The text output produced by the model.

        Raises:
            ValueError: If required arguments are missing or the request payload is empty.
            RuntimeError: If the Grok response does not contain any textual content.
        """
        if not api_key:
            raise ValueError("api_key must be provided.")
        if not prompt and not messages and not images:
            raise ValueError("At least one of prompt, messages, or images must be provided.")
        if not model:
            raise ValueError("model must be provided.")

        normalized_messages = normalize_messages(prompt=prompt, messages=messages)
        message_parts: list[object] = []
        if normalized_messages:
            for message in normalized_messages:
                if message["content"]:
                    message_parts.append(f"{message['role']}: {message['content']}")

        if images:
            for image in images:
                message_parts.append(chat_image(self._to_image_url(image)))

        if not message_parts:
            raise ValueError("No content provided for response generation.")

        retry_count = clamp_retries(max_retries)

        def _run_request() -> str:
            grok_client = Client(api_key=api_key)

            create_kwargs = {
                "model": model,
                "max_tokens": max_tokens,
                "messages": [user(*message_parts)],
            }

            normalized_effort = (reasoning_effort or "").strip().lower()
            if normalized_effort in {"low", "high"}:
                create_kwargs["reasoning_effort"] = normalized_effort
            if temperature is not None:
                create_kwargs["temperature"] = float(temperature)
            if top_p is not None:
                create_kwargs["top_p"] = float(top_p)

            if timeout_s is not None:
                create_kwargs["timeout"] = timeout_s

            try:
                chat = grok_client.chat.create(**create_kwargs)
                response = chat.sample()
            except Exception as exc:
                logger.exception(
                    "xAI Grok chat request failed: %s request_id=%s",
                    exc,
                    request_id,
                )
                raise

            content = getattr(response, "content", None)
            if content:
                logger.info(
                    "xAI chat succeeded: model=%s images=%d text_len=%d request_id=%s",
                    model,
                    len(images or []),
                    len(content or ""),
                    request_id,
                )
                return content

            reasoning_content = getattr(response, "reasoning_content", None)
            if reasoning_content:
                logger.info(
                    "xAI chat succeeded (reasoning): model=%s images=%d text_len=%d request_id=%s",
                    model,
                    len(images or []),
                    len(reasoning_content or ""),
                    request_id,
                )
                return reasoning_content

            # Treat successful calls without textual content as a successful, empty response
            # rather than raising. This aligns with callers that handle empty outputs gracefully.
            logger.info(
                "xAI chat succeeded with no text: model=%s images=%d request_id=%s",
                model,
                len(images or []),
                request_id,
            )
            return ""

        return run_with_retries(
            func=_run_request,
            max_retries=retry_count,
            retry_backoff_s=retry_backoff_s,
            request_id=request_id,
        )

    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:
        return await run_sync_in_thread(
            lambda: self.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,
        *,
        api_key: str,
        prompt: str,
        model: str,
        image_size: str = "2K",
        image: Optional[ImageInput] = None,
    ) -> bytes:
        """Generate an image using the Grok API.

        Raises:
            NotImplementedError: This method is not yet implemented for Grok.
        """
        raise NotImplementedError("Image generation is not implemented for Grok.")

    async def async_generate_image(
        self,
        *,
        api_key: str,
        prompt: str,
        model: str,
        image_size: str = "2K",
        image: Optional[ImageInput] = None,
    ) -> bytes:
        return await run_sync_in_thread(
            lambda: self.generate_image(
                api_key=api_key,
                prompt=prompt,
                model=model,
                image_size=image_size,
                image=image,
            )
        )

    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,
    ) -> list[dict[str, Optional[str]]]:
        """Return the Grok language models available to the authenticated account."""
        if not api_key:
            raise ValueError("api_key must be provided.")

        retry_count = clamp_retries(max_retries)

        def _run_request() -> list[dict[str, Optional[str]]]:
            grok_client = Client(api_key=api_key)
            models: list[dict[str, Optional[str]]] = []

            try:
                iterator = grok_client.models.list_language_models()
            except Exception as exc:
                logger.exception(
                    "xAI list language models failed: %s request_id=%s",
                    exc,
                    request_id,
                )
                raise

            for model in iterator:
                model_id = getattr(model, "name", None)
                if not model_id:
                    continue

                aliases = getattr(model, "aliases", None)
                display_name: Optional[str] = None
                if aliases:
                    try:
                        display_name = next(iter(aliases)) or None
                    except (StopIteration, TypeError):
                        display_name = None

                models.append(
                    {
                        "id": model_id,
                        "display_name": display_name,
                    }
                )

            logger.info(
                "xAI list_language_models succeeded: count=%d request_id=%s",
                len(models),
                request_id,
            )
            return models

        return run_with_retries(
            func=_run_request,
            max_retries=retry_count,
            retry_backoff_s=retry_backoff_s,
            request_id=request_id,
        )

    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,
    ) -> list[dict[str, Optional[str]]]:
        return await run_sync_in_thread(
            lambda: self.list_models(
                api_key=api_key,
                request_id=request_id,
                timeout_s=timeout_s,
                max_retries=max_retries,
                retry_backoff_s=retry_backoff_s,
            )
        )

    @staticmethod
    def _to_image_url(image: ImageInput) -> str:
        """Convert an image reference into a URL or data URL suitable for the xAI SDK."""
        if isinstance(image, Path):
            return _encode_image_path(image)

        if image.startswith(("http://", "https://", "data:")):
            return image

        return _encode_image_path(Path(image))


def _encode_image_path(path: Path) -> str:
    """Generate a data URL for the provided image path."""
    expanded = path.expanduser()
    data = expanded.read_bytes()
    mime_type = mimetypes.guess_type(expanded.name)[0] or "application/octet-stream"
    encoded = base64.b64encode(data).decode("utf-8")
    return f"data:{mime_type};base64,{encoded}"
