modules/google_image_gen.py

import os
import base64
from loguru import logger
from typing import Optional, Dict, Any

# Tenta importar o SDK do Google GenAI
try:
    from google import genai
    from google.genai import types
    HAS_GENAI = True
except ImportError:
    HAS_GENAI = False

class GoogleImageGenerator:
    """
    Gerador de imagens usando Google Imagen 3 (Nano Banana).
    """
    def __init__(self, api_key: Optional[str] = None):
        self.api_key = api_key or os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
        self.client = None
        if HAS_GENAI and self.api_key:
            try:
                # Tenta inicializar o cliente sem forçar versão para deixar o SDK decidir
                self.client = genai.Client(api_key=self.api_key)
                # Verifica versão do SDK
                try:
                    import google.genai as genai_mod
                    logger.info(f"✅ Google GenAI SDK v{getattr(genai_mod, '__version__', 'unknown')} inicializado")
                except:
                    logger.info("✅ Google Image Generator (Imagen 3) inicializado")
            except Exception as e:
                logger.error(f"❌ Falha ao inicializar Google GenAI Client: {e}")

    def generate(self, prompt: str, aspect_ratio: str = "1:1", model: str = "flux") -> Dict[str, Any]:
        """
        Gera uma imagem via Pollinations como primário e Imagen 3 como fallback.
        """
        # ✅ PRIORIDADE: Pollinations (Poly/Flux) é o preferido do usuário agora
        try:
            res = self._pollinations_fallback(prompt, aspect_ratio, model)
            if res.get("success"):
                return res
        except Exception as poly_err:
            logger.warning(f"⚠️ Pollinations falhou: {poly_err}. Tentando Google como fallback final...")

        # 🔄 FALLBACK: Google Imagen (Apenas se o Pollinations falhar)
        if not self.client:
            return {"success": False, "error": "Google GenAI Client não disponível ou sem chave API"}

        try:
            # Mapeamento de aspect ratio para o formato do Imagen
            # Imagen 3 suporta: "1:1", "4:3", "3:4", "16:9", "9:16"
            valid_ratios = ["1:1", "4:3", "3:4", "16:9", "9:16"]
            if aspect_ratio not in valid_ratios:
                aspect_ratio = "1:1"

            logger.info(f"🎨 Gerando imagem via Imagen 3 (Nano Banana): '{prompt[:50]}...' [{aspect_ratio}]")
            
            # Debug: Listar modelos disponíveis E FILTRAR válidos
            available_models = []
            try:
                available_models = [m.name.replace("models/", "") for m in self.client.models.list()]
                logger.info(f"📋 Modelos disponíveis: {available_models[:10]}...")  # Trunca log
            except Exception as le:
                logger.warning(f"Não foi possível listar modelos: {le}")

            # Modelos preferidos (env override first)
            preferred_models = os.getenv("GEMINI_IMAGE_MODEL", "").split(",") if os.getenv("GEMINI_IMAGE_MODEL") else []
            models_to_try = preferred_models + [
                "imagen-3.0-generate-001",
                "imagen-3.0-fast-001", 
                "imagen-4.0-generate-001",
                "nano-banana-pro-preview"
            ]
            
            # FILTRA APENAS MODELOS QUE REALMENTE EXISTEM
            models_to_try = [m.strip() for m in models_to_try if m.strip() in available_models]
            if not models_to_try:
                logger.error("❌ Nenhum modelo Imagen válido disponível!")
                return self._pollinations_fallback(prompt, aspect_ratio, model)
            
            logger.info(f"🎨 Tentando modelos válidos: {models_to_try[:3]}...")
            
            last_err = None
            for model_id in models_to_try:
                # Tenta tanto singular quanto plural (o SDK mudou entre versões beta/GA)
                for method_name in ["generate_image", "generate_images"]:
                    if not hasattr(self.client.models, method_name):
                        continue
                        
                    try:
                        logger.debug(f"Tentando {method_name} com {model_id}...")
                        method = getattr(self.client.models, method_name)
                        
                        # Ajusta a classe de config conforme o método
                        config_class = types.GenerateImageConfig if method_name == "generate_image" else types.GenerateImagesConfig
                        
                        img_response = method(
                            model=model_id,
                            prompt=prompt,
                            config=config_class(
                                number_of_images=1,
                                aspect_ratio=aspect_ratio,
                            )
                        )
                        
                        if img_response and img_response.generated_images:
                            img = img_response.generated_images[0]
                            logger.success(f"✅ Imagem gerada com sucesso via Google ({model_id})!")
                            # Verifica todos os possíveis locais do buffer
                            data = None
                            if hasattr(img, 'image') and hasattr(img.image, 'data'): data = img.image.data
                            elif hasattr(img, 'image') and hasattr(img.image, 'image_bytes'): data = img.image.image_bytes
                            elif hasattr(img, 'image_bytes'): data = img.image_bytes
                            
                            if data:
                                return {
                                    "success": True, 
                                    "buffer": data,
                                    "mime_type": "image/png",
                                    "model": model_id
                                }
                    except Exception as inner_e:
                        last_err = inner_e
                        logger.warning(f"⚠️ {method_name} com {model_id} falhou: {inner_e}")
                        continue
            # Se chegou aqui, todos os modelos do Google falharam
            logger.warning(f"⚠️ Todos os modelos Imagen falharam. Erro: {last_err}")
            
            # Fallback para Pollinations (Gratuito)
            return self._pollinations_fallback(prompt, aspect_ratio, model)

        except Exception as e:
            logger.error(f"❌ Erro crítico no gerador de imagem: {e}")
            return {"success": False, "error": str(e)}


    def _pollinations_fallback(self, prompt: str, aspect_ratio: str = "1:1", model: str = "flux") -> Dict[str, Any]:
        """
        Gera uma imagem via Pollinations.ai como fallback gratuito.
        """
        try:
            import requests
            import urllib.parse
            import random
            
            logger.info(f"🎙️ Usando Pollinations ({model}) para: '{prompt[:50]}...'")
            
            # Dimensões baseadas no aspect ratio
            width, height = 1024, 1024
            if aspect_ratio == "16:9": width, height = 1280, 720
            elif aspect_ratio == "9:16": width, height = 720, 1280
            elif aspect_ratio == "4:3": width, height = 1024, 768
            elif aspect_ratio == "3:4": width, height = 768, 1024
            
            seed = random.randint(1, 999999)
            
            encoded_prompt = urllib.parse.quote(prompt)
            url = f"https://image.pollinations.ai/prompt/{encoded_prompt}?width={width}&height={height}&model={model}&seed={seed}&nologo=true"
            
            response = requests.get(url, timeout=45)
            if response.status_code == 200 and len(response.content) > 1000:
                logger.success(f"✅ Imagem gerada com sucesso via Pollinations ({model})!")
                return {
                    "success": True,
                    "buffer": response.content,
                    "mime_type": "image/png",
                    "model": f"pollinations-{model}"
                }
            
            return {"success": False, "error": f"Pollinations falhou com status {response.status_code}"}
        except Exception as e:
            logger.error(f"❌ Erro no fallback Pollinations: {e}")
            return {"success": False, "error": f"Erro no fallback: {str(e)}"}

# Singleton
_instance = None
def get_google_image_gen():
    global _instance
    if _instance is None:
        _instance = GoogleImageGenerator()
    return _instance