"""Thin wrapper around the Google Gemini API via the google-genai SDK."""

from __future__ import annotations

import base64
import mimetypes
from pathlib import Path
import logging
from typing import Optional, Sequence, Union
from urllib.error import URLError
from urllib.request import urlopen

from google import genai
from google.genai import types

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


class GeminiClient:
    """Convenience wrapper around the Google Gemini SDK."""

    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 Gemini model.

        Args:
            api_key: API key used to authenticate with the Gemini API.
            prompt: Natural-language instruction or query for the model.
            model: Identifier of the Gemini model to target (for example, ``"gemini-2.5-flash"``).
            max_tokens: Cap for tokens across the entire exchange, defaults to 32000.
            reasoning_effort: Included for API parity; currently unused by the Gemini SDK.
            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.
            URLError: If an image URL cannot be retrieved.
            google.genai.errors.APIError: If the underlying Gemini request fails.
        """
        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.")

        parts: list[types.Part] = []
        if prompt:
            parts.append(types.Part.from_text(text=prompt))

        if images:
            for image in images:
                parts.append(self._to_image_part(image))

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

        content = types.Content(role="user", parts=parts)

        config = types.GenerateContentConfig(max_output_tokens=max_tokens)
        # reasoning_effort is accepted for compatibility but not currently applied because the
        # Gemini SDK does not expose an equivalent configuration parameter.

        client = genai.Client(api_key=api_key)
        try:
            try:
                response = client.models.generate_content(
                    model=model,
                    contents=[content],
                    config=config,
                )
            except Exception as exc:
                logger.exception("Gemini generate_content failed: %s", exc)
                raise
        finally:
            closer = getattr(client, "close", None)
            if callable(closer):
                try:
                    closer()
                except Exception:
                    pass

        if response.text:
            result_text = response.text
            logger.info(
                "Gemini generate_content succeeded: model=%s images=%d text_len=%d",
                model,
                len(images or []),
                len(result_text or ""),
            )
            return result_text

        candidate_texts: list[str] = []
        for candidate in getattr(response, "candidates", []) or []:
            content_obj = getattr(candidate, "content", None)
            if not content_obj:
                continue
            for part in getattr(content_obj, "parts", []) or []:
                text = getattr(part, "text", None)
                if text:
                    candidate_texts.append(text)

        if candidate_texts:
            result_text = "\n".join(candidate_texts)
            logger.info(
                "Gemini generate_content succeeded (candidates): model=%s images=%d text_len=%d",
                model,
                len(images or []),
                len(result_text or ""),
            )
            return result_text

        # 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(
            "Gemini generate_content 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 models available to the authenticated Gemini account."""
        if not api_key:
            raise ValueError("api_key must be provided.")

        models: list[dict[str, Optional[str]]] = []
        client = genai.Client(api_key=api_key)
        try:
            try:
                iterator = client.models.list()
            except Exception as exc:
                logger.exception("Gemini list models failed: %s", exc)
                raise
            for model in iterator:
                model_id = getattr(model, "name", None)
                if model_id is None and isinstance(model, dict):
                    model_id = model.get("name")
                if not model_id:
                    continue

                # Normalize IDs like "models/<id>" -> "<id>"
                if isinstance(model_id, str) and model_id.startswith("models/"):
                    model_id = model_id.split("/", 1)[1]

                display_name = getattr(model, "display_name", None)
                if display_name is None and isinstance(model, dict):
                    display_name = model.get("display_name")

                models.append({"id": model_id, "display_name": display_name})
        finally:
            closer = getattr(client, "close", None)
            if callable(closer):
                try:
                    closer()
                except Exception:
                    pass

        logger.info("Gemini list_models succeeded: count=%d", len(models))
        return models

    @staticmethod
    def _to_image_part(image: ImageInput) -> types.Part:
        """Convert an image reference into a Gemini SDK part."""
        if isinstance(image, Path):
            return _part_from_path(image)

        if image.startswith("data:"):
            return _part_from_data_url(image)

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

        return _part_from_path(Path(image))


def _part_from_path(path: Path) -> types.Part:
    """Create an image part from a local filesystem path."""
    expanded = path.expanduser()
    data = expanded.read_bytes()
    mime_type = mimetypes.guess_type(expanded.name)[0] or "application/octet-stream"
    return types.Part.from_bytes(data=data, mime_type=mime_type)


def _part_from_url(url: str) -> types.Part:
    """Create an image part by downloading content from a URL."""
    with urlopen(url) as response:
        data = response.read()
        mime_type = response.info().get_content_type()

    if not mime_type or mime_type == "application/octet-stream":
        mime_type = mimetypes.guess_type(url)[0] or "application/octet-stream"

    return types.Part.from_bytes(data=data, mime_type=mime_type)


def _part_from_data_url(data_url: str) -> types.Part:
    """Create an image part from a data URL."""
    header, encoded = data_url.split(",", 1)
    metadata = header[len("data:") :]
    mime_type = "application/octet-stream"

    if ";" in metadata:
        mime_type, _, metadata = metadata.partition(";")
    elif metadata:
        mime_type = metadata

    if "base64" in metadata:
        data = base64.b64decode(encoded)
    else:
        data = encoded.encode("utf-8")

    return types.Part.from_bytes(data=data, mime_type=mime_type or "application/octet-stream")
