import os
from huggingface_hub import InferenceClient
import gradio as gr
from PIL import Image, ImageDraw

# === CONFIGURACIÓN DEL TOKEN ===
HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("HF_TOKEN no encontrado. Configúralo en los Secrets del Space.")

client = InferenceClient(token=HF_TOKEN)

# === MODELO ÚNICO: FLUX.1-schnell (público, sin gate) ===
MODELO = "black-forest-labs/FLUX.1-schnell"

# === FUNCIÓN DE GENERACIÓN CON MANEJO DE ERRORES ===
def generar_imagen(prompt):
    negative_prompt = "low quality, blurry, cartoon, lowres, deformed, mutation, extra limbs, bad anatomy, distorted, unnatural colors, pixelated, watermark, text, logo"

    try:
        print(f"Generando imagen con FLUX.1-schnell... Prompt: {prompt[:50]}...")

        image = client.text_to_image(
            model=MODELO,
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=576,      # 9:16
            height=1024,    # 9:16
            guidance_scale=3.5,
            num_inference_steps=4  # schnell es rápido
        )
        return image

    except Exception as e:
        print(f"Error en generación: {e}")
        # Imagen de error roja
        img = Image.new("RGB", (576, 1024), "#8B0000")
        draw = ImageDraw.Draw(img)
        try:
            draw.text((20, 480), f"Error: {str(e)[:80]}", fill="white", font_size=30)
        except:
            draw.text((20, 480), "Error al generar imagen", fill="white")
        return img

# === INTERFAZ GRADIO ===
with gr.Blocks(title="FLUX 9:16 Generator") as demo:
    gr.Markdown("# Generador de Imágenes 9:16 con FLUX.1-schnell")
    gr.Markdown("**Formato: 576x1024 px (ideal para móviles)**")
    gr.Markdown("_Negative prompt automático: evita baja calidad, deformaciones, texto, etc._")

    with gr.Row():
        with gr.Column(scale=2):
            prompt_input = gr.Textbox(
                label="Describe tu imagen (en inglés para mejores resultados)",
                placeholder="Ej: A futuristic city at sunset, cyberpunk style, vibrant colors, detailed, 9:16",
                lines=4
            )
            btn = gr.Button("Generar Imagen", variant="primary", size="lg")

        with gr.Column(scale=1):
            output_image = gr.Image(
                label="Imagen Generada",
                type="pil",
                height=600
            )

    # Botón → función
    btn.click(fn=generar_imagen, inputs=prompt_input, outputs=output_image)

    # Ejemplos
    gr.Examples(
        examples=[
            ["A serene Japanese garden with cherry blossoms, koi pond, morning mist, ultra-realistic, 9:16"],
            ["Cyberpunk girl with neon hair, rain, reflections on wet street, cinematic lighting, 9:16"],
            ["Minimalist product shot of a luxury watch on black marble, studio lighting, 9:16"]
        ],
        inputs=prompt_input
    )

# === LANZAMIENTO ===
if __name__ == "__main__":
    demo.launch()