import base64
import json
from datetime import datetime
from pathlib import Path
import google.oauth2.service_account as sa
import vertexai
from vertexai.preview.vision_models import ImageGenerationModel


class GeminiImageChat:
    """Chat interactivo para generación de imágenes con Gemini"""

    def __init__(self, credentials_path=None):
        """Inicializa el cliente de Gemini"""
        # Buscar archivo de credenciales
        if credentials_path is None:
            credentials_path = Path(__file__).parent.parent / "bnt-ia-innovacion-new.json"

        credentials_path = Path(credentials_path)

        if not credentials_path.exists():
            raise ValueError(f"Archivo de credenciales no encontrado: {credentials_path}")

        # Leer credenciales del archivo JSON
        with open(credentials_path, 'r') as f:
            credentials_json = json.load(f)

        # Cargar credenciales usando el método correcto
        credentials = sa.Credentials.from_service_account_file(str(credentials_path))

        # Inicializar Vertex AI
        vertexai.init(
            project=credentials_json.get('project_id'),
            credentials=credentials,
            location='us-central1'
        )

        # Inicializar modelo de generación de imágenes
        self.model = ImageGenerationModel.from_pretrained("gemini-2.5-flash-image")
        self.output_dir = Path("generated_images")
        self.output_dir.mkdir(exist_ok=True)

    def save_image(self, image_data, prompt):
        """Guarda la imagen generada en disco"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        # Sanitizar el prompt para usar como parte del nombre
        safe_prompt = "".join(c for c in prompt[:30] if c.isalnum() or c in (' ', '_')).rstrip()
        safe_prompt = safe_prompt.replace(' ', '_')

        filename = f"{timestamp}_{safe_prompt}.png"
        filepath = self.output_dir / filename

        # Decodificar y guardar la imagen
        if isinstance(image_data, str):
            # Si es base64
            image_bytes = base64.b64decode(image_data)
        else:
            image_bytes = image_data

        with open(filepath, 'wb') as f:
            f.write(image_bytes)

        return filepath

    def generate_image(self, prompt, aspect_ratio="1:1", save=True):
        """Genera una imagen basada en el prompt"""
        print(f"\n{'='*60}")
        print(f"Generando imagen para: {prompt}")
        print(f"{'='*60}\n")

        try:
            # Generar imágenes usando Vertex AI
            images = self.model.generate_images(
                prompt=prompt,
                number_of_images=1,
                language="es",
                aspect_ratio=aspect_ratio,
                safety_filter_level="block_some",
                person_generation="allow_adult",
            )

            # Guardar imágenes generadas
            saved_files = []
            if images and save:
                for idx, image in enumerate(images.images):
                    # Convertir imagen a bytes
                    image_bytes = image._image_bytes

                    # Guardar usando el método existente
                    filepath = self.save_image(image_bytes, prompt)
                    saved_files.append(filepath)
                    print(f"✓ Imagen {idx + 1} guardada en: {filepath}")

            return {
                "images": saved_files,
                "success": True
            }

        except Exception as e:
            print(f"\n✗ Error al generar imagen: {str(e)}")
            return {
                "images": [],
                "success": False,
                "error": str(e)
            }

    def start_chat(self):
        """Inicia el chat interactivo"""
        print("\n" + "="*60)
        print("🎨 CHAT DE GENERACIÓN DE IMÁGENES CON GEMINI")
        print("="*60)
        print("\nComandos disponibles:")
        print("  - Escribe un prompt para generar una imagen")
        print("  - 'config' - Cambiar configuración de imagen")
        print("  - 'salir' - Salir del chat")
        print("="*60 + "\n")

        aspect_ratio = "1:1"

        while True:
            try:
                user_input = input("\n💭 Prompt: ").strip()

                if not user_input:
                    continue

                if user_input.lower() == 'salir':
                    print("\n👋 ¡Hasta pronto!")
                    break

                elif user_input.lower() == 'config':
                    print("\nConfiguración actual:")
                    print(f"  Aspect ratio: {aspect_ratio}")
                    print("\nOpciones de aspect ratio: 1:1, 16:9, 9:16, 3:4, 4:3")

                    new_ratio = input("\nNuevo aspect ratio (Enter para mantener): ").strip()
                    if new_ratio:
                        aspect_ratio = new_ratio

                    print(f"\n✓ Configuración actualizada: {aspect_ratio}")

                else:
                    # Generar imagen
                    self.generate_image(user_input, aspect_ratio)

            except KeyboardInterrupt:
                print("\n\n👋 ¡Hasta pronto!")
                break
            except Exception as e:
                print(f"\n✗ Error: {str(e)}")


def main():
    """Función principal"""
    try:
        chat = GeminiImageChat()
        chat.start_chat()
    except ValueError as e:
        print(f"✗ Error de configuración: {e}")
        print("\nAsegúrate de tener el archivo bnt-ia-innovacion-new.json en el directorio raíz")
    except Exception as e:
        print(f"✗ Error inesperado: {e}")


if __name__ == "__main__":
    main()