import gradio as gr
import requests
import io
import random
import os
from PIL import Image

# Configuración global
API_TOKEN = os.getenv("HF_READ_TOKEN")
HEADERS = {"Authorization": f"Bearer {API_TOKEN}"}
TIMEOUT = 120

# Opciones de modelos LORA
LORA_OPTIONS = {
    "Aitana": "RedSparkie/aitaoca2",
    "Alice": "RedSparkie/wonderland", 
    "Andrea": "RedSparkie/bote",
    "Halle": "RedSparkie/quema",
    "Lluna": "RedSparkie/luna",
    "Lola": "RedSparkie/lola",
    "Maxi": "RedSparkie/maximo", 
    "Pruckute": "RedSparkie/pruckute",
    "Rivers": "RedSparkie/larios4"
}

DEFAULT_MODEL = "black-forest-labs/FLUX.1-dev"

# Función principal
def query(lora_name, prompt, steps, cfg_scale, randomize_seed, seed, width, height):
    if not prompt:
        return None
    
    lora_id = LORA_OPTIONS.get(lora_name, DEFAULT_MODEL)
    api_url = f"https://api-inference.huggingface.co/models/{lora_id.strip()}"
    
    if randomize_seed:
        seed = random.randint(1, 4294967296)
    
    # Payload
    payload = {
        "inputs": f"{prompt}, ultra detail, ultra elaboration, ultra quality, perfect.",
        "steps": steps,
        "cfg_scale": cfg_scale,
        "seed": seed,
        "width": width,
        "height": height
    }
    
    # Petición con reintentos
    for attempt in range(3):  # Máximo de 3 intentos
        response = requests.post(api_url, headers=HEADERS, json=payload, timeout=TIMEOUT)
        if response.status_code == 200:
            try:
                image = Image.open(io.BytesIO(response.content))
                return image, str(seed)
            except Exception as e:
                raise gr.Error(f"Error al procesar la imagen: {e}")
        elif response.status_code == 503:
            if attempt < 2:  # Si no es el último intento
                time.sleep(5)
            else:
                raise gr.Error("El modelo está cargándose, por favor intenta más tarde.")
        else:
            raise gr.Error(f"Error inesperado: {response.status_code} - {response.text}")
    return None

# Interfaz
with gr.Blocks() as app:
    gr.Markdown("<center><h1>Generador de Imágenes</h1></center>")
    
    with gr.Row():
        lora_selector = gr.Dropdown(label="Modelo", choices=list(LORA_OPTIONS.keys()), value="", allow_custom_value=True)
        text_prompt = gr.Textbox(label="Prompt", placeholder="Describe tu imagen")
    
    with gr.Row():
        width = gr.Slider(label="Ancho", minimum=64, maximum=1280, step=8, value=1024)
        height = gr.Slider(label="Alto", minimum=64, maximum=1280, step=8, value=1024)
        steps = gr.Slider(label="Pasos", minimum=1, maximum=100, step=1, value=20)
        cfg = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
        randomize_seed = gr.Checkbox(label="Randomizar semilla", value=True)
        seed = gr.Textbox(label="Semilla", placeholder="-1 para aleatorio", value="-1")

    image_output = gr.Image(type="pil", label="Resultado")
    seed_output = gr.Textbox(label="Semilla utilizada", interactive=False)
    generate_button = gr.Button("Generar")

    generate_button.click(
        query,
        inputs=[lora_selector, text_prompt, steps, cfg, randomize_seed, seed, width, height],
        outputs=[image_output, seed_output]
    )

app.launch(share=True)