import os
import gradio as gr
import requests
from PIL import Image
from io import BytesIO
from openai import OpenAI
from dotenv import load_dotenv
import time
import base64
import random

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

class FluxNineteenGenerator:
    """Generatore di immagini utilizzando l'API FLUX di Nineteen.ai"""
    
    def __init__(self):
        """Inizializza il generatore"""
        print("Inizializzazione del generatore di immagini FLUX Nineteen.ai...")
        
        # Ottieni l'API key da variabile d'ambiente
        self.api_key = os.getenv("NINETEEN_API_KEY")
        if not self.api_key:
            raise ValueError(
                "È necessario configurare la variabile d'ambiente NINETEEN_API_KEY.\n"
                "Su Hugging Face Spaces: Aggiungi la variabile nelle impostazioni dello Space.\n"
                "In locale: Configura la variabile nel tuo ambiente di sviluppo."
            )
        
        # Configurazione dell'API
        self.api_endpoint = "https://api.nineteen.ai/v1/text-to-image"
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        # Modello predefinito
        self.default_model = "black-forest-labs/FLUX.1-schnell"
        print(f"Generatore inizializzato con modello predefinito: {self.default_model}")
        
        # Lista dei modelli disponibili
        self.models = {
            "schnell": "black-forest-labs/FLUX.1-schnell",
            "proteus": "dataautogpt3/ProteusV0.4-Lightning",
            "dreamshaper": "Lykon/dreamshaper-xl-lightning"
        }
    
    def generate_image(self, prompt, model=None, steps=28, cfg_scale=7.0, height=1024, width=1024, negative_prompt=None):
        """
        Genera un'immagine utilizzando l'API FLUX.
        
        Args:
            prompt (str): Descrizione dell'immagine da generare
            model (str, optional): Nome del modello da utilizzare. Default: None (usa il modello predefinito)
            steps (int, optional): Numero di passi di inferenza. Default: 28
            cfg_scale (float, optional): Guidance scale. Default: 7.0
            height (int, optional): Altezza dell'immagine. Default: 1024
            width (int, optional): Larghezza dell'immagine. Default: 1024
            negative_prompt (str, optional): Prompt negativo. Default: None
        
        Returns:
            tuple: (PIL.Image, str) L'immagine generata e un messaggio di stato
        """
        try:
            # Usa il modello predefinito se non specificato
            if not model:
                model = self.default_model
            
            # Se viene fornito un nome breve, usa il nome completo dal dizionario
            if model in self.models:
                model = self.models[model]
            
            # Prepara i parametri per la richiesta
            payload = {
                "prompt": prompt,
                "model": model,
                "num_inference_steps": steps,
                "guidance_scale": cfg_scale,
                "height": height,
                "width": width
            }
            
            if negative_prompt:
                payload["negative_prompt"] = negative_prompt
            
            print(f"Generazione immagine con prompt: '{prompt}'")
            print(f"Parametri: modello={model}, steps={steps}, cfg_scale={cfg_scale}, dimensioni={height}x{width}")
            
            start_time = time.time()
            
            # Invia la richiesta all'API
            response = requests.post(
                self.api_endpoint,
                headers=self.headers,
                json=payload
            )
            
            # Gestisci gli errori HTTP
            if response.status_code != 200:
                error_message = f"Errore API: {response.status_code} - {response.text}"
                print(error_message)
                return None, error_message
            
            # Decodifica l'immagine
            end_time = time.time()
            try:
                image_b64 = response.json()["image_b64"]
                image_data = base64.b64decode(image_b64)
                image = Image.open(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 Exception as e:
                error_message = f"Errore nella decodifica dell'immagine: {str(e)}"
                print(error_message)
                return None, error_message
                
        except Exception as e:
            error_message = f"Errore durante la generazione: {str(e)}"
            print(error_message)
            return None, error_message

class GeminiPromptGenerator:
    """Generatore di prompt creativi utilizzando Gemini AI"""
    
    def __init__(self):
        """Inizializza il generatore di prompt"""
        self.client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=os.getenv("OPENROUTER_API_KEY")
        )
        self.model = {
            "model": "google/gemma-3-27b-it:free",
            "temperature": 1.4,      # Aumentata per più creatività
            "top_p": 0.99,          # Aumentato per più variabilità
            "max_tokens": 8192,      # Equivalente a max_output_tokens
            "extra_headers": {
                "HTTP-Referer": "https://your-site.com",  # Richiesto da OpenRouter
                "X-Title": "AI Image Generator"           # Richiesto da OpenRouter
            }
        }
        
        # Lista di stili artistici per variare i prompt
        self.art_styles = [
            "hyperrealistic", "cinematic", "digital art", "oil painting",
            "concept art", "fantasy art", "sci-fi", "surrealism",
            "studio photography", "anime", "cyberpunk", "steampunk"
        ]
        
        # Lista di atmosfere per variare i prompt
        self.moods = [
            "ethereal", "mysterious", "dramatic", "peaceful",
            "energetic", "melancholic", "whimsical", "epic",
            "intimate", "nostalgic", "futuristic", "dreamy"
        ]
        
    def generate_creative_prompt(self):
        """Genera un prompt creativo per la generazione di immagini"""
        # Seleziona casualmente uno stile e un'atmosfera
        style = random.choice(self.art_styles)
        mood = random.choice(self.moods)
        
        system_prompt = f"""You are an expert prompt engineer for artistic image generation.
        Create a unique and creative prompt to generate an extraordinary image.
        Current time: {time.strftime('%Y-%m-%d %H:%M:%S')}
        
        The prompt must:
        - Be in English
        - Create a completely new and original scene, different from any previous ones
        - Focus on the suggested style: {style}
        - Incorporate the suggested mood: {mood}
        - Include specific details about:
          * Main subject and its unique characteristics
          * Rich environmental details and setting
          * Dynamic lighting and atmospheric effects
          * Specific color palette and visual elements
          * Camera angle or perspective
        - Use vivid and evocative language
        - Be about 2-3 lines long
        - Include elements that work well with the FLUX.1-schnell model
        - End with style keywords: {style}, {mood}, photorealistic, cinematic, 8K, ultra-detailed
        
        Generate ONLY the prompt, without any explanations or introductions.
        IMPORTANT: Create something completely different from previous prompts.
        DO NOT include words like 'prompt' or 'description' in the response."""
        
        try:
            response = self.client.chat.completions.create(
                messages=[{"role": "user", "content": system_prompt}],
                **self.model
            )
            if response.choices[0].message.content:
                return response.choices[0].message.content.strip()
            else:
                return "Sorry, I couldn't generate a creative prompt. Please try again."
        except Exception as e:
            return f"Error generating prompt: {str(e)}"

def create_ui():
    """Crea l'interfaccia utente Gradio"""
    
    # Inizializza i generatori
    flux_gen = FluxNineteenGenerator()
    gemini_gen = GeminiPromptGenerator()
    
    def generate_random_prompt():
        """Genera un prompt casuale usando Gemini"""
        try:
            return gemini_gen.generate_creative_prompt()
        except Exception as e:
            return f"Error: {str(e)}"
    
    def generate_image(prompt, model, steps, cfg_scale, height, width, negative_prompt):
        """Funzione per generare immagini dall'interfaccia"""
        try:
            return flux_gen.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, str(e)
    
    with gr.Blocks(title="FLUX Image Generator") as demo:
        gr.Markdown("# 🎨 FLUX Image Generator")
        gr.Markdown("Image generator with FLUX.1-schnell and Gemini AI")
        
        with gr.Row():
            with gr.Column(scale=1):
                prompt_input = gr.Textbox(
                    label="Prompt",
                    placeholder="Enter your image description or click 'Generate Creative Prompt'...",
                    lines=3
                )
                
                with gr.Row():
                    generate_prompt_btn = gr.Button("🤖 Generate Creative Prompt", variant="secondary")
                    clear_prompt_btn = gr.Button("🗑️ Clear", variant="secondary")
                
                with gr.Accordion("Advanced Settings", open=False):
                    model_input = gr.Dropdown(
                        choices=list(FluxNineteenGenerator().models.keys()),
                        value="schnell",
                        label="Model"
                    )
                    
                    steps_input = gr.Slider(
                        minimum=1,
                        maximum=100,
                        value=28,
                        step=1,
                        label="Steps"
                    )
                    
                    cfg_input = gr.Slider(
                        minimum=1.0,
                        maximum=20.0,
                        value=7.0,
                        step=0.5,
                        label="CFG Scale"
                    )
                    
                    with gr.Row():
                        height_input = gr.Slider(
                            minimum=512,
                            maximum=1536,
                            value=1024,
                            step=64,
                            label="Height (px)"
                        )
                        width_input = gr.Slider(
                            minimum=512,
                            maximum=1536,
                            value=1024,
                            step=64,
                            label="Width (px)"
                        )
                    
                    negative_prompt_input = gr.Textbox(
                        label="Negative Prompt",
                        placeholder="Elements to avoid in the image...",
                        lines=2
                    )
                
                generate_image_btn = gr.Button("🎨 Generate Image", variant="primary")
            
            with gr.Column(scale=1):
                output_image = gr.Image(label="Generated Image")
                output_status = gr.Textbox(label="Status", interactive=False)
        
        # Eventi
        generate_prompt_btn.click(
            fn=generate_random_prompt,
            outputs=prompt_input
        )
        
        clear_prompt_btn.click(
            fn=lambda: "",
            outputs=prompt_input
        )
        
        generate_image_btn.click(
            fn=generate_image,
            inputs=[
                prompt_input,
                model_input,
                steps_input,
                cfg_input,
                height_input,
                width_input,
                negative_prompt_input
            ],
            outputs=[output_image, output_status]
        )
    
    return demo

if __name__ == "__main__":
    demo = create_ui()
    demo.launch()