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._format_prompt(input_data.prompt, input_data.negative_prompt)
        original_prompt = prompt

        if input_data.enhance_prompt:
            prompt = self._enhance_prompt(prompt)

        enhanced_prompt = prompt if input_data.enhance_prompt else original_prompt

        try:
            # Nano Banana: generate_content with response_modalities
            config = genai.types.GenerateContentConfig(
                response_modalities=["IMAGE"],
            )

            response = self.client.models.generate_content(
                model=self.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 candidate.content and candidate.content.parts:
                        for part in candidate.content.parts:
                            if hasattr(part, "inline_data") and part.inline_data:
                                generated_images.append(  # noqa: PERF401
                                    self._create_generated_image_data(
                                        part.inline_data.data, content_type
                                    )
                                )

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

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

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