"""
Gemini API client for Gemini 2.5 Flash Image generation.
Uses the generateContent API endpoint per Google's documentation.
"""

import base64
import logging
from typing import Any

import httpx

from ..config.constants import GEMINI_API_BASE, GEMINI_MODELS
from ..core.exceptions import (
    APIError,
    AuthenticationError,
    ContentPolicyError,
    RateLimitError,
)

logger = logging.getLogger(__name__)


class GeminiClient:
    """Client for Gemini 2.5 Flash Image API."""

    def __init__(self, api_key: str, timeout: int = 60):
        """
        Initialize Gemini client.

        Args:
            api_key: Gemini API key
            timeout: Request timeout in seconds
        """
        self.api_key = api_key
        self.timeout = timeout
        self.base_url = GEMINI_API_BASE
        self.client = httpx.AsyncClient(timeout=timeout)

    async def generate_image(
        self,
        prompt: str,
        *,
        model: str = "gemini-2.5-flash-image",
        input_image: str | None = None,
        aspect_ratio: str | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """
        Generate or edit an image using Gemini 2.5 Flash Image.

        Args:
            prompt: Text prompt for image generation or editing instruction
            model: Model to use (default: gemini-2.5-flash-image)
            input_image: Base64-encoded input image for editing (optional)
            aspect_ratio: Desired aspect ratio (optional, influences output)
            **kwargs: Additional parameters

        Returns:
            Dict with 'images' key containing list of base64-encoded image data

        Raises:
            APIError: If the API request fails
        """
        model_id = GEMINI_MODELS.get(model, model)
        url = f"{self.base_url}/models/{model_id}:generateContent"

        # Build request body according to doc.md
        parts: list[dict[str, Any]] = []

        # Add input image if provided (for editing)
        if input_image:
            parts.append({
                "inline_data": {
                    "mime_type": "image/png",
                    "data": input_image
                }
            })

        # Add text prompt (include aspect ratio hint if specified)
        prompt_text = prompt
        if aspect_ratio:
            prompt_text = f"{prompt}. Aspect ratio: {aspect_ratio}"

        parts.append({"text": prompt_text})

        request_body = {"contents": [{"parts": parts}]}

        headers = {
            "x-goog-api-key": self.api_key,
            "Content-Type": "application/json",
        }

        try:
            logger.debug(f"Sending request to {url}")
            response = await self.client.post(url, json=request_body, headers=headers)
            response.raise_for_status()
            data = response.json()

            # Extract images from response
            images = self._extract_images(data)

            if not images:
                raise APIError("No image data found in Gemini API response")

            return {"images": images, "model": model, "response": data}

        except httpx.HTTPStatusError as e:
            self._handle_http_error(e)
        except Exception as e:
            logger.error(f"Gemini API request failed: {e}")
            raise APIError(f"Gemini API request failed: {e}")

    async def generate_text(
        self,
        prompt: str,
        *,
        model: str = "gemini-2.0-flash",
        system_instruction: str | None = None,
    ) -> str:
        """
        Generate text using Gemini (for prompt enhancement).

        Args:
            prompt: Text prompt
            model: Model to use
            system_instruction: Optional system instruction

        Returns:
            Generated text response
        """
        model_id = GEMINI_MODELS.get(model, model)
        url = f"{self.base_url}/models/{model_id}:generateContent"

        request_body = {
            "contents": [{"parts": [{"text": prompt}]}]
        }

        if system_instruction:
            request_body["system_instruction"] = {
                "parts": [{"text": system_instruction}]
            }

        headers = {
            "x-goog-api-key": self.api_key,
            "Content-Type": "application/json",
        }

        try:
            response = await self.client.post(url, json=request_body, headers=headers)
            response.raise_for_status()
            data = response.json()

            # Extract text from response
            text = self._extract_text(data)
            return text

        except httpx.HTTPStatusError as e:
            self._handle_http_error(e)
        except Exception as e:
            logger.error(f"Gemini text generation failed: {e}")
            raise APIError(f"Gemini text generation failed: {e}")

    def _extract_images(self, response_data: dict[str, Any]) -> list[str]:
        """Extract base64 image data from Gemini API response."""
        images = []

        try:
            candidates = response_data.get("candidates", [])
            for candidate in candidates:
                content = candidate.get("content", {})
                parts = content.get("parts", [])
                for part in parts:
                    if "inline_data" in part:
                        image_data = part["inline_data"].get("data")
                        if image_data:
                            images.append(image_data)
        except Exception as e:
            logger.warning(f"Error extracting images from response: {e}")

        return images

    def _extract_text(self, response_data: dict[str, Any]) -> str:
        """Extract text from Gemini API response."""
        try:
            candidates = response_data.get("candidates", [])
            if not candidates:
                return ""

            content = candidates[0].get("content", {})
            parts = content.get("parts", [])

            # Concatenate all text parts
            text_parts = [part.get("text", "") for part in parts if "text" in part]
            return "".join(text_parts)

        except Exception as e:
            logger.warning(f"Error extracting text from response: {e}")
            return ""

    def _handle_http_error(self, error: httpx.HTTPStatusError) -> None:
        """Handle HTTP errors and raise appropriate exceptions."""
        status_code = error.response.status_code
        error_text = error.response.text

        logger.error(f"API request failed with status {status_code}: {error_text}")

        if status_code == 401 or status_code == 403:
            raise AuthenticationError(
                "Authentication failed. Please check your Gemini API key.",
                status_code=status_code
            )
        elif status_code == 429:
            raise RateLimitError(
                "Rate limit exceeded. Please try again later.",
                status_code=status_code
            )
        elif status_code == 400 and ("SAFETY" in error_text.upper() or "BLOCKED" in error_text.upper()):
            raise ContentPolicyError(
                "Content was blocked by safety filters. Please modify your prompt.",
                status_code=status_code
            )
        else:
            raise APIError(
                f"API request failed with status {status_code}: {error_text}",
                status_code=status_code
            )

    async def close(self) -> None:
        """Close the HTTP client."""
        await self.client.aclose()
