"""
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"""
    with gr.Blocks(title="FLUX Image Generator") as interface:
        gr.Markdown("# 🎨 FLUX Image Generator")
        gr.Markdown("Genera immagini creative utilizzando il modello FLUX.1-schnell di Nineteen.ai")
        
        with gr.Row():
            with gr.Column(scale=4):
                prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="Descrivi l'immagine che vuoi generare...",
                    lines=3
                )
                
                with gr.Row():
                    submit_btn = gr.Button("🎨 Genera", variant="primary")
                    clear_btn = gr.Button("🗑️ Pulisci")
            
            with gr.Column(scale=6):
                image_output = gr.Image(label="Immagine Generata", type="filepath")
                
        # Eventi
        submit_btn.click(
            fn=generator.generate_image,
            inputs=prompt,
            outputs=image_output
        )
        clear_btn.click(
            fn=lambda: (None, ""),
            inputs=None,
            outputs=[image_output, prompt]
        )
        
        return interface

def main():
    """Funzione principale"""
    # Crea il generatore
    generator = FluxNineteenGenerator()
    
    # Crea l'interfaccia
    interface = create_ui(generator)
    
    # Ottieni le credenziali dalle variabili d'ambiente
    username = os.getenv("GRADIO_USERNAME")
    password = os.getenv("GRADIO_PASSWORD")
    
    # Parametri di base per il lancio
    launch_kwargs = {
        "server_name": "0.0.0.0",
    }
    
    # Verifica le credenziali e configura l'autenticazione se necessario
    if username and password:
        print(f"Autenticazione configurata con username: {username}")
        launch_kwargs["auth"] = [(username, password)]
    else:
        print("Autenticazione disabilitata. Su HF Spaces, imposta GRADIO_USERNAME e GRADIO_PASSWORD.")
    
    # Rileva se siamo su Hugging Face Spaces
    if os.getenv("SPACE_ID") is None:
        # Non siamo su HF Spaces, aggiungi share per renderlo accessibile
        launch_kwargs["share"] = True
    
    # Avvia l'interfaccia
    interface.launch(**launch_kwargs)

if __name__ == "__main__":
    main()