"""
FLUX Image Generator - Nineteen.ai API (Versione HF Spaces)
--------------------------------------
Questo script utilizza l'API di Nineteen.ai per generare immagini
con il modello FLUX.1-schnell, offrendo un'interfaccia grafica intuitiva.
Versione ottimizzata per Hugging Face Spaces con autenticazione.

Autore: Utente
Data: 09/03/2025
"""

import os
import io
import time
import base64
import requests
import gradio as gr
from PIL import Image
from dotenv import load_dotenv

# Carica le variabili d'ambiente (se presenti)
load_dotenv()

class FluxNineteenGenerator:
    """
    Classe per generare immagini utilizzando il modello FLUX.1-schnell
    attraverso l'API di Nineteen.ai.
    """
    
    def __init__(self):
        """
        Inizializza il generatore di immagini FLUX tramite API Nineteen.ai.
        """
        print("Inizializzazione del generatore di immagini FLUX Nineteen.ai...")
        
        # Ottieni l'API key da variabile d'ambiente o usa quella predefinita
        self.api_key = os.getenv("NINETEEN_API_KEY", "rayon_6nzcBRzaq8f7iRzFUX2IsscMBGKgmJcO")
        
        # Configurazione dell'API
        self.api_endpoint = "https://api.nineteen.ai/v1/text-to-image"
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "accept": "application/json",
            "Content-Type": "application/json"
        }
        
        # Modelli disponibili
        self.models = ["black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-dev"]
        self.default_model = "black-forest-labs/FLUX.1-schnell"
        
        print(f"Generatore inizializzato con modello predefinito: {self.default_model}")
    
    def generate_image(self, prompt, model=None, steps=8, cfg_scale=3, 
                      height=1024, width=1024, negative_prompt=""):
        """
        Genera un'immagine utilizzando l'API di Nineteen.ai.
        
        Args:
            prompt (str): Descrizione testuale dell'immagine da generare
            model (str, optional): Modello da utilizzare
            steps (int): Numero di passi di inferenza
            cfg_scale (float): Scala di guidance per la generazione
            height (int): Altezza dell'immagine in pixel
            width (int): Larghezza dell'immagine in pixel
            negative_prompt (str): Prompt negativo per escludere elementi indesiderati
            
        Returns:
            PIL.Image: L'immagine generata
            str: Messaggio di stato
        """
        if not model:
            model = self.default_model
            
        try:
            # Prepara il payload
            data = {
                "prompt": prompt,
                "model": model,
                "steps": steps,
                "cfg_scale": cfg_scale,
                "height": height,
                "width": width,
                "negativePrompt": negative_prompt
            }
            
            # Log della richiesta
            print(f"Generazione immagine con prompt: '{prompt}'")
            print(f"Parametri: modello={model}, steps={steps}, cfg_scale={cfg_scale}, dimensioni={width}x{height}")
            
            # Effettua la chiamata API
            start_time = time.time()
            response = requests.post(self.api_endpoint, headers=self.headers, json=data)
            end_time = time.time()
            
            # Gestione degli errori
            if response.status_code != 200:
                error_message = f"Errore API: {response.status_code} - {response.text}"
                print(error_message)
                return None, error_message
            
            # Estrai l'immagine in formato base64 dalla risposta
            try:
                image_b64 = response.json()["image_b64"]
                image_data = base64.b64decode(image_b64)
                image = Image.open(io.BytesIO(image_data))
                
                print(f"Immagine generata in {end_time - start_time:.2f} secondi")
                return image, f"Immagine generata in {end_time - start_time:.2f} secondi"
            except KeyError:
                return None, f"Errore: La risposta API non contiene il campo 'image_b64'. Risposta: {response.json()}"
            except Exception as e:
                return None, f"Errore nel decodificare l'immagine: {str(e)}"
                
        except Exception as e:
            error_message = f"Errore durante la generazione dell'immagine: {str(e)}"
            print(error_message)
            return None, error_message

def create_ui(generator):
    """
    Crea l'interfaccia utente Gradio.
    
    Args:
        generator (FluxNineteenGenerator): Istanza del generatore
    
    Returns:
        gradio.Interface: L'interfaccia Gradio
    """
    
    def generate_image_ui(prompt, model, steps, cfg_scale, height, width, negative_prompt):
        """Funzione per generare immagini dall'interfaccia"""
        # Validazione dei parametri
        if not prompt or len(prompt.strip()) == 0:
            return None, "Il prompt non può essere vuoto."
        
        try:
            # Converti i parametri al tipo corretto
            steps = int(steps)
            cfg_scale = float(cfg_scale)
            height = int(height)
            width = int(width)
            
            # Validazione dei valori
            if steps < 1 or steps > 50:
                return None, "Il numero di passi deve essere compreso tra 1 e 50."
            if cfg_scale < 1 or cfg_scale > 30:
                return None, "La guidance scale deve essere compresa tra 1 e 30."
            if height < 512 or height > 1536:
                return None, "L'altezza deve essere compresa tra 512 e 1536 pixel."
            if width < 512 or width > 1536:
                return None, "La larghezza deve essere compresa tra 512 e 1536 pixel."
            
            # Genera l'immagine
            return generator.generate_image(
                prompt=prompt,
                model=model,
                steps=steps,
                cfg_scale=cfg_scale,
                height=height,
                width=width,
                negative_prompt=negative_prompt
            )
        except Exception as e:
            return None, f"Errore: {str(e)}"
    
    # Crea i componenti dell'interfaccia
    with gr.Blocks(title="FLUX Nineteen.ai Image Generator") as interface:
        gr.Markdown("# FLUX Nineteen.ai Image Generator")
        gr.Markdown("Genera immagini utilizzando il modello FLUX.1-schnell tramite l'API di Nineteen.ai")
        
        with gr.Row():
            with gr.Column(scale=3):
                prompt_input = gr.Textbox(
                    label="Prompt",
                    placeholder="Descrivi l'immagine che desideri generare...",
                    lines=3
                )
                
                negative_prompt_input = gr.Textbox(
                    label="Prompt Negativo (opzionale)",
                    placeholder="Elementi da escludere dall'immagine...",
                    lines=2
                )
                
                model_input = gr.Dropdown(
                    generator.models,
                    label="Modello",
                    value=generator.default_model
                )
                
                with gr.Row():
                    steps_input = gr.Slider(
                        minimum=1,
                        maximum=50,
                        value=8,
                        step=1,
                        label="Passi di Inferenza"
                    )
                    cfg_input = gr.Slider(
                        minimum=1.0,
                        maximum=30.0,
                        value=3.0,
                        step=0.1,
                        label="Guidance Scale (CFG)"
                    )
                
                with gr.Row():
                    height_input = gr.Slider(
                        minimum=512,
                        maximum=1536,
                        value=1024,
                        step=64,
                        label="Altezza (px)"
                    )
                    width_input = gr.Slider(
                        minimum=512,
                        maximum=1536,
                        value=1024,
                        step=64,
                        label="Larghezza (px)"
                    )
                
                generate_button = gr.Button("Genera Immagine", variant="primary")
            
            with gr.Column(scale=4):
                output_image = gr.Image(label="Immagine Generata", type="pil")
                output_status = gr.Textbox(label="Stato", interactive=False)
        
        # Esempi di prompt
        with gr.Accordion("Esempi di Prompt", open=False):
            gr.Markdown("""
            ### Esempi di prompt ottimizzati per FLUX.1-schnell
            
            Clicca su uno degli esempi per utilizzarlo:
            """)
            
            examples = [
                ["A breathtaking view of the Dolomites at sunrise, golden light illuminating the jagged peaks, morning mist rising from the valley below, ultra-detailed, cinematic, 8K resolution, photorealistic"],
                ["Futuristic Tokyo cityscape at night, neon lights reflecting on wet streets after rain, towering skyscrapers with holographic advertisements, flying vehicles, photorealistic, cinematic lighting, 8K"],
                ["Portrait of a weathered old fisherman with deep wrinkles and piercing blue eyes, wearing a cable-knit sweater, salt and pepper beard, golden hour lighting, ultra-detailed skin texture, photorealistic"],
                ["Massive space station orbiting Jupiter, with Earth visible in the distance, detailed mechanical structures, solar panels, docking bays with spacecraft, photorealistic, NASA quality, 8K"],
                ["Bioluminescent forest at night with giant mushrooms, glowing plants, mystical atmosphere, small magical creatures, ultra-detailed vegetation, photorealistic textures, fantasy world with realistic lighting"]
            ]
            
            gr.Examples(
                examples=examples,
                inputs=prompt_input
            )
        
        # Collega il pulsante di generazione
        generate_button.click(
            generate_image_ui,
            inputs=[
                prompt_input, 
                model_input, 
                steps_input, 
                cfg_input, 
                height_input, 
                width_input, 
                negative_prompt_input
            ],
            outputs=[output_image, output_status]
        )
    
    return interface

# Funzione principale
def main():
    """Funzione principale"""
    # Crea il generatore
    generator = FluxNineteenGenerator()
    
    # Crea l'interfaccia
    interface = create_ui(generator)
    
    # Configura l'autenticazione
    username = os.getenv("GRADIO_USERNAME")
    password = os.getenv("GRADIO_PASSWORD")
    
    # Verifica se le credenziali sono disponibili
    if username and password:
        auth = (username, password)
        print(f"Autenticazione configurata con username: {username}")
    else:
        # Per test locali, disabilitiamo l'autenticazione
        auth = None
        print("Autenticazione disabilitata per test locali. Su HF Spaces, imposta GRADIO_USERNAME e GRADIO_PASSWORD.")
    
    # Avvia l'interfaccia con o senza autenticazione
    interface.launch(auth=auth)

if __name__ == "__main__":
    main()