# BATUTO-ART MIX - Código completo corregido
# (FLUX / SD1.5 / REVE CREATE con API key editable)
# Optimizado para CPU en HuggingFace Spaces

import os
import gradio as gr
from diffusers import DiffusionPipeline
import torch
import requests
from PIL import Image
from io import BytesIO

# ==============================
#   CONFIGURACIÓN BASE CPU
# ==============================
DEVICE = "cpu"
torch.set_grad_enabled(False)

def load_flux(model_id):
    pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
    pipe.to(DEVICE)
    pipe.enable_attention_slicing()
    return pipe

# Cache de modelos
MODEL_CACHE = {}

# ==============================
#   GENERADOR FLUX
# ==============================
def generate_flux(model_name, prompt, steps, guidance, width, height, seed):
    if model_name not in MODEL_CACHE:
        MODEL_CACHE[model_name] = load_flux(model_name)
    pipe = MODEL_CACHE[model_name]

    generator = torch.manual_seed(seed) if seed else None

    image = pipe(
        prompt=prompt,
        num_inference_steps=steps,
        guidance_scale=guidance,
        width=width,
        height=height,
        generator=generator
    ).images[0]

    out = "/tmp/flux_output.png"
    image.save(out)
    return out

# ==============================
#   GENERADOR SD1.5
# ==============================
def load_sd15():
    pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float32)
    pipe.to(DEVICE)
    pipe.enable_attention_slicing()
    return pipe

# SD15 load único
def generate_sd(prompt, steps, guidance, width, height, seed):
    if "sd15" not in MODEL_CACHE:
        MODEL_CACHE["sd15"] = load_sd15()
    pipe = MODEL_CACHE["sd15"]

    generator = torch.manual_seed(seed) if seed else None

    image = pipe(
        prompt=prompt,
        num_inference_steps=steps,
        guidance_scale=guidance,
        width=width,
        height=height,
        generator=generator
    ).images[0]

    out = "/tmp/sd15_output.png"
    image.save(out)
    return out

# ==============================
#   REVE CREATE
# ==============================
def reve_generate(prompt, key, model):
    if not key:
        return None

    url = "https://api.reveai.xyz/v1/images"
    headers = {"Authorization": f"Bearer {key}"}
    data = {"prompt": prompt, "model": model}

    resp = requests.post(url, json=data, headers=headers)

    if resp.status_code != 200:
        return None

    img_url = resp.json().get("image")
    img_data = requests.get(img_url).content
    img = Image.open(BytesIO(img_data))

    out = "/tmp/reve.png"
    img.save(out)
    return out

# ==============================
#   UI COMPLETA
# ==============================
def build_ui():
    with gr.Blocks(title="BATUTO-ART MIX") as demo:

        gr.Markdown("# 🖼️ **BATUTO-ART MIX** ")

        with gr.Tabs():

            # ============================
            # TAB: FLUX
            # ============================
            with gr.Tab("FLUX.2 / 1-Schnell"):
                flux_prompt = gr.Textbox(label="Prompt", lines=3)
                model_select = gr.Dropdown([
                    "black-forest-labs/FLUX.1-schnell",
                    "black-forest-labs/FLUX.1-dev",
                    "black-forest-labs/FLUX.2-dev"
                ], value="black-forest-labs/FLUX.1-schnell", label="Modelo FLUX")

                steps = gr.Slider(5, 50, value=20, label="Steps")
                guidance = gr.Slider(0, 10, value=3, label="Guidance Scale")
                seed = gr.Number(value=0, label="Seed (0 = aleatorio)")

                width = gr.Number(value=576, label="Width")
                height = gr.Number(value=1024, label="Height 9:16")

                btn_flux = gr.Button("Generar Imagen")
                out_flux_img = gr.Image(label="Resultado")
                out_flux_file = gr.File(label="Descargar archivo")

                # Acción
                btn_flux.click(
                    fn=lambda m,p,s,g,w,h,sd: generate_flux(m,p,int(s),float(g),int(w),int(h),int(sd)),
                    inputs=[model_select, flux_prompt, steps, guidance, width, height, seed],
                    outputs=[out_flux_file]
                )

                # Mostrar imagen automáticamente
                out_flux_file.change(fn=lambda f: Image.open(f) if f else None, inputs=[out_flux_file], outputs=[out_flux_img])

            # ============================
            # TAB: SD1.5
            # ============================
            with gr.Tab("Stable Diffusion 1.5"):
                sd_prompt = gr.Textbox(label="Prompt", lines=3)
                sd_steps = gr.Slider(5, 50, value=20)
                sd_guidance = gr.Slider(0, 10, value=3)
                sd_seed = gr.Number(value=0)
                sd_width = gr.Number(value=576)
                sd_height = gr.Number(value=1024)

                btn_sd = gr.Button("Generar Imagen")
                out_sd_img = gr.Image(label="Resultado")
                out_sd_file = gr.File(label="Descargar archivo")

                btn_sd.click(
                    fn=lambda p,s,g,w,h,sd: generate_sd(p,int(s),float(g),int(w),int(h),int(sd)),
                    inputs=[sd_prompt, sd_steps, sd_guidance, sd_width, sd_height, sd_seed],
                    outputs=[out_sd_file]
                )

                out_sd_file.change(fn=lambda f: Image.open(f) if f else None, inputs=[out_sd_file], outputs=[out_sd_img])

            # ============================
            # TAB: REVE CREATE
            # ============================
            with gr.Tab("REVE CREATE"):
                reve_api = gr.Textbox(label="API Key REVE")
                reve_prompt = gr.Textbox(label="Prompt", lines=3)
                reve_model = gr.Dropdown([
                    "reve-1",
                    "reve-2",
                    "reve-fast"
                ], value="reve-fast", label="Modelo REVE")

                btn_reve = gr.Button("Generar Imagen")
                reve_out_img = gr.Image(label="Resultado")

                reve_btn_download = gr.Button("Descargar Imagen", variant="primary")
                reve_out_file = gr.File(label="Archivo generado")

                last_file = gr.State()

                btn_reve.click(
                    fn=lambda p,k,m: reve_generate(p,k,m),
                    inputs=[reve_prompt, reve_api, reve_model],
                    outputs=[last_file]
                )

                last_file.change(fn=lambda f: Image.open(f) if f else None, inputs=[last_file], outputs=[reve_out_img])

                reve_btn_download.click(
                    fn=lambda f: f,
                    inputs=[last_file],
                    outputs=[reve_out_file]
                )

    return demo

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