import logging

from google import genai

from appkit_imagecreator.backend.generators.google import GoogleImageGenerator
from appkit_imagecreator.backend.models import (
    GeneratedImageData,
    GenerationInput,
    ImageGeneratorResponse,
    ImageResponseState,
)

logger = logging.getLogger(__name__)


class NanoBananaImageGenerator(GoogleImageGenerator):
    """Generator for Google's Nano Banana image models.

    Uses generate_content API with response_modalities=['IMAGE'].
    Supports:
    - Nano Banana: gemini-2.0-flash-exp
    - Nano Banana alternatives: gemini-exp-1206
    """

    async def _perform_generation(
        self, input_data: GenerationInput
    ) -> ImageGeneratorResponse:
        """Generate images using Nano Banana via generate_content API."""
        prompt = self._enhance_prompt(input_data)

        try:
            config_params = {"response_modalities": ["IMAGE"]}
            image_config_params = {}

            # Known fields for ImageConfig
            image_config_fields = {
                "aspect_ratio",
                "image_size",
                "person_generation",
                "output_mime_type",
                "output_compression_quality",
            }

            if self.model.config:
                for key, value in self.model.config.items():
                    if key in image_config_fields:
                        image_config_params[key] = value
                    else:
                        config_params[key] = value

            # Use requested aspect ratio if not forced by model config
            if "aspect_ratio" not in image_config_params:
                image_config_params["aspect_ratio"] = self._aspect_ratio(
                    input_data.width, input_data.height
                )

            if image_config_params:
                config_params["image_config"] = genai.types.ImageConfig(
                    **image_config_params
                )

            config = genai.types.GenerateContentConfig(**config_params)

            response = self.client.models.generate_content(
                model=self.model.model,
                contents=prompt,
                config=config,
            )

            output_format = "png"
            content_type = f"image/{output_format}"
            generated_images: list[GeneratedImageData] = []

            # Extract images from response
            if response.candidates:
                for candidate in response.candidates:
                    if not (candidate.content and candidate.content.parts):
                        continue

                    generated_images.extend(
                        self._create_generated_image_data(
                            part.inline_data.data, content_type
                        )
                        for part in candidate.content.parts
                        if hasattr(part, "inline_data") and part.inline_data
                    )

            if not generated_images:
                logger.warning(
                    "No images generated by Nano Banana model %s", self.model.id
                )
                return ImageGeneratorResponse(
                    state=ImageResponseState.FAILED,
                    generated_images=[],
                    error="No images returned from API",
                    enhanced_prompt=prompt,
                )

            return ImageGeneratorResponse(
                state=ImageResponseState.SUCCEEDED,
                generated_images=generated_images,
                enhanced_prompt=prompt,
            )

        except Exception as e:
            logger.exception(
                "Error generating image with %s model %s",
                self.model.id,
                self.model.model,
            )
            return ImageGeneratorResponse(
                state=ImageResponseState.FAILED,
                generated_images=[],
                error=str(e),
                enhanced_prompt=prompt,
            )
