"""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, Union

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

ImageInput = Union[str, Path]
logger = logging.getLogger(__name__)


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

    def generate_response(
        self,
        *,
        api_key: str,
        prompt: str,
        model: str,
        max_tokens: int = 32000,
        reasoning_effort: Optional[str] = None,
        images: Optional[Sequence[ImageInput]] = None,
    ) -> 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"``).
            images: Optional collection of image references (local paths, URLs, or data URLs).

        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 images:
            raise ValueError("At least one of prompt or images must be provided.")
        if not model:
            raise ValueError("model must be provided.")

        message_parts = []
        if prompt:
            message_parts.append(prompt)

        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.")

        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

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

        content = getattr(response, "content", None)
        if content:
            logger.info(
                "xAI chat succeeded: model=%s images=%d text_len=%d",
                model,
                len(images or []),
                len(content or ""),
            )
            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",
                model,
                len(images or []),
                len(reasoning_content or ""),
            )
            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",
            model,
            len(images or []),
        )
        return ""

    def list_models(self, *, api_key: str) -> 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.")

        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", exc)
            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", len(models))
        return models

    @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}"
