import os

# Descargas grandes (modelo base ~51GB): usar hf_transfer (multi-conexión, resumible)
# y subir el timeout de lectura para tolerar cortes de red transitorios.
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "120")

MOCK = False #not os.environ.get("HF_SPACE_ID")

if (MOCK):
    import sys
    from unittest.mock import MagicMock

    # Creamos un módulo falso llamado 'spaces'
    mock_spaces = MagicMock()

    # Definimos el decorador GPU para que simplemente devuelva la función original sin cambios
    def mock_gpu_decorator(duration=None):
        def decorator(func):
            return func
        return decorator

    mock_spaces.GPU = mock_gpu_decorator

    # Lo insertamos en los módulos del sistema para que 'import spaces' funcione
    sys.modules["spaces"] = mock_spaces


import gradio as gr
import json
import logging
from PIL import Image
import spaces
import copy
import random
import time
import re
import math
import numpy as np
import traceback
import tempfile
from prompt_rewrite import rewrite
from gradio_client import Client, handle_file
import hashlib
from functools import partial

if (not MOCK):
    import torch
    from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
    from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download



NUM_LORAS = 4
base_model = "Qwen/Qwen-Image-2512"

def apply_aspect_ratio(ratio):
    sizes = {
        "1:1": (1024, 1024),
        "16:9": (1365, 768),
        "9:16": (768, 1365),
        "3:2": (1254, 836),
        "2:3": (836, 1254),
        "3:1": (1774, 591),
        "2:1": (1448, 724),
    }
    return sizes.get(ratio, (1024, 1024))

DEFAULT_ASPECT_RATIO = "16:9"



# ✅ NUEVO: importar optimización avanzada tipo Qwen-Image-MultipleAngles
#from optimization import optimize_pipeline_

LORAS_CACHE = {
    "data": [],
    "last_hash": None,
}
def load_loras_hot():
    if MOCK:
        return load_loras_from_file()

    """Load loras.json and detect changes."""
    path = hf_hub_download(
        repo_id="lichorosario/qwen-image-lora-dlc-v3",
        filename="loras.json",
        repo_type="space",
    )

    with open(path, "r", encoding="utf-8") as f:
        raw = f.read()

    current_hash = hashlib.sha256(raw.encode("utf-8")).hexdigest()

    if current_hash != LORAS_CACHE["last_hash"]:
        LORAS_CACHE["data"] = json.loads(raw)
        LORAS_CACHE["last_hash"] = current_hash
        print("🔁 LoRA config updated")

    return LORAS_CACHE["data"]

    
# Load LoRAs from JSON file
def load_loras_from_file():
    """Load LoRA configurations from external JSON file."""
    try:
        with open('loras.json', 'r', encoding='utf-8') as f:
            return json.load(f)
    except FileNotFoundError:
        print("Warning: loras.json file not found. Using empty list.")
        return []
    except json.JSONDecodeError as e:
        print(f"Error parsing loras.json: {e}")
        return []


# Load the LoRAs
#loras = load_loras_from_file()
loras = load_loras_hot()

# loras = load_loras_hot()
loras = load_loras_hot()

selected_loras = []

if not MOCK:
    # Initialize the base model
    dtype = torch.bfloat16
    device = "cuda" if torch.cuda.is_available() else "cpu"



# Scheduler configuration from the Qwen-Image-Lightning repository
scheduler_config = {
    "base_image_seq_len": 256,
    "base_shift": math.log(3),
    "invert_sigmas": False,
    "max_image_seq_len": 8192,
    "max_shift": math.log(3),
    "num_train_timesteps": 1000,
    "shift": 1.0,
    "shift_terminal": None,
    "stochastic_sampling": False,
    "time_shift_type": "exponential",
    "use_beta_sigmas": False,
    "use_dynamic_shifting": True,
    "use_exponential_sigmas": False,
    "use_karras_sigmas": False,
}


if not MOCK:
    scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)

    pipe = DiffusionPipeline.from_pretrained(
        "Qwen/Qwen-Image-2512", scheduler=scheduler, torch_dtype=dtype
    ).to(device)







# Lightning LoRA info (no global state)
LIGHTNING_LORA_REPO = "lightx2v/Qwen-Image-2512-Lightning"
LIGHTNING_LORA_WEIGHT = "Qwen-Image-2512-Lightning-4steps-V1.0-fp32.safetensors"
LIGHTNING8_LORA_WEIGHT = "Qwen-Image-2512-Lightning-8steps-V1.0-fp32.safetensors"
LIGHTNING_FP8_4STEPS_LORA_WEIGHT = "Qwen-Image-fp8-e4m3fn-Lightning-4steps-V1.0-bf16.safetensors"

WULI_LORA_REPO = "Wuli-art/Qwen-Image-2512-Turbo-LoRA"
WULI_LORA_WEIGHT = "Wuli-Qwen-Image-2512-Turbo-LoRA-4steps-V3.0-bf16.safetensors"

WULI2STP_LORA_REPO = "Wuli-art/Qwen-Image-2512-Turbo-LoRA-2-Steps"
WULI2STP_LORA_WEIGHT = "Wuli-Qwen-Image-2512-Turbo-LoRA-2steps-V1.0-bf16.safetensors"


MAX_SEED = np.iinfo(np.int32).max


### MODIFICACIÓN 1: AÑADIR FUNCIONES PARA GESTIONAR EL HISTORIAL ###
def update_history(new_images, history):
    """Añade las nuevas imágenes generadas al principio de la lista del historial."""
    if history is None:
        history = []
    if new_images is not None and len(new_images) > 0:
        updated_history = new_images + history
        return updated_history[:24]
    return history


def clear_history():
    """Devuelve una lista vacía para limpiar la galería de historial."""
    return []


PID_UPSCALER_SPACE = "prithivMLmods/PiD-Image-Upscaler"


def select_history_image(history, evt: gr.SelectData):
    """Guarda la imagen del historial que se acaba de clickear."""
    if not history or evt.index is None or evt.index >= len(history):
        return None
    return history[evt.index][0]


def upscale_history_image(image, oauth_token: gr.OAuthToken | None = None):
    """Escala 4x la imagen elegida del historial llamando a la API del space
    prithivMLmods/PiD-Image-Upscaler: el modelo corre en el hardware de ese
    space, no en el nuestro. Con el token OAuth del visitante la cuota ZeroGPU
    de esa llamada se le factura a él; sin login sale anónima."""
    if image is None:
        raise gr.Error("Click an image in the History gallery first.")

    input_path = os.path.join(tempfile.mkdtemp(), "upscale_input.png")
    image.save(input_path)

    try:
        client = Client(PID_UPSCALER_SPACE, hf_token=getattr(oauth_token, "token", None))
        _, slider = client.predict(handle_file(input_path), "", api_name="/upscaler_run")
    except Exception as e:
        raise gr.Error(f"PiD upscaler failed: {e}")

    # El endpoint remoto devuelve el gr.update(...) sin resolver: un dict
    # {"value": [original, upscalada]}, no la tupla ya lista.
    original, upscaled = slider["value"]
    return gr.update(visible=True, value=(original, upscaled))
### FIN DE LA MODIFICACIÓN 1 ###




class calculateDuration:
    def __init__(self, activity_name=""):
        self.activity_name = activity_name


    def __enter__(self):
        self.start_time = time.time()
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        self.end_time = time.time()
        self.elapsed_time = self.end_time - self.start_time
        if self.activity_name:
            print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
        else:
            print(f"Elapsed time: {self.elapsed_time:.6f} seconds")




def update_selection(evt: gr.SelectData, width, height):
    selected_lora = loras[evt.index]
    versions = selected_lora.get("versions", [])
    if versions:
        # Use first version or logic to pick one? Logic in get_selection handles the "active" one.
        # This function updates just the placeholder and info text, might not be fully used anymore or overlapping logic.
        # We will keep it but maybe it's not the primary update mechanism for state.
        pass
        
    new_placeholder = f"Type a prompt for {selected_lora['title']}"
    lora_repo = selected_lora["repo"]
    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨"
    
    examples_list = []
    try:
        model_card = ModelCard.load(lora_repo)
        widget_data = model_card.data.get("widget", [])
        if widget_data and len(widget_data) > 0:
            for example in widget_data[:4]:
                if "output" in example and "url" in example["output"]:
                    image_url = f"https://huggingface.co/{lora_repo}/resolve/main/{example['output']['url']}"
                    prompt_text = example.get("text", "")
                    examples_list.append([prompt_text])
    except Exception as e:
        print(f"Could not load model card for {lora_repo}: {e}")
    
    return (
        gr.update(placeholder=new_placeholder),
        updated_text,
        evt.index,
        width,
        width,
        height
    )


def update_version_selection(version_title, current_state, idx):
    """
    Handle version dropdown change.
    Update state and the specific Markdown to show the version's image.
    """
    # current_state is list of tuples: (image_index, version_index, scale)
    # idx is the column index (0-based)
    
    new_state = [tuple(item) for item in current_state]
    current_image_index = new_state[idx][0]
    
    if current_image_index is None:
        return gr.update(), new_state # Should not happen if logic is correct

    # Handle legacy state (2 elements) or new state (3 elements)
    if len(new_state[idx]) == 3:
        current_image_index, current_version_index, current_scale = new_state[idx]
    else:
        current_image_index, current_scale = new_state[idx]
        current_version_index = None

    lora = loras[current_image_index]
    versions = lora.get("versions", [])
    
    # Find selected version object
    selected_version = next((v for i, v in enumerate(versions) if v["title"] == version_title), None)
    selected_version_index = next((i for i, v in enumerate(versions) if v["title"] == version_title), None)
    
    if selected_version:
        # Update Markdown image
        lora_name = lora["title"]
        version_image = selected_version.get("image", lora.get("image"))
        markdown_content = f"<img src='{version_image}' style='height:100px; display:block; margin-bottom:10px;' />\n\n**{lora_name}** ({version_title})"
        
        
        # Update state: (image_index, version_index, scale)
        new_state[idx] = (current_image_index, selected_version_index, current_scale)
        
        return gr.update(value=markdown_content), new_state
    
    return gr.update(), new_state


def handle_speed_mode(speed_mode):
    """Update UI based on speed/quality toggle."""
    if speed_mode == "light 4":
        return gr.update(value="Light mode (4 steps) selected"), 4, 1.0
    elif speed_mode == "light 4 fp8":
        return gr.update(value="Light mode (4 steps fp8) selected"), 4, 1.0
    elif speed_mode == "light 8":
        return gr.update(value="Light mode (8 steps) selected"), 8, 1.0

    elif speed_mode == "Wuli-art":
        return gr.update(value="Light mode (4 steps) Wuli-art selected"), 4, 1.0

    elif speed_mode == "Wuli2Step-art":
        return gr.update(value="Light mode (2 steps) Wuli-art selected"), 2, 1.0

    
    else: 
        return gr.update(value="Normal quality (45 steps) selected"), 30, 3.5




def gpu_duration(*args, **kwargs):
    """Presupuesto de GPU dinámico: pasos × imágenes × resolución, con margen.

    La cuota se cobra por tiempo real usado; esto solo ajusta la reserva
    (mejor prioridad de cola y menos daño si un task cuelga). Ante cualquier
    duda cae al default de 120s.
    """
    try:
        prompts = [p for p in args[0:4] if p and str(p).strip()]
        n_prompts = max(1, len(prompts))
        steps = int(args[6])
        width, height = int(args[10]), int(args[11])
        mult = float(str(args[13]).replace("x", ""))
        real_quantity = int(args[14]) + 1
        n_images = n_prompts * real_quantity
        pixels_factor = max(1.0, (width * mult * height * mult) / (1024 * 1024))
        per_image = 3.0 * steps * pixels_factor
        return int(min(900, 60 + per_image * n_images))
    except Exception:
        return 120


@spaces.GPU(duration=gpu_duration, size="xlarge")
def run_lora_multi(
    prompt_1, prompt_2, prompt_3, prompt_4,
    negative_prompt,
    cfg_scale, steps,
    selected_loras_state,  # Changed from selected_index
    randomize_seed, seed,
    width, height,
    speed_mode,
    quality_multiplier,
    quantity,
    history,
    prompt_enhance=False,
    progress=gr.Progress(track_tqdm=True)
):
    # selected_loras_state is a list of tuples: [(image_index, version_index, scale), ...]
    if selected_loras_state is None:
        selected_loras_state = []

    # Filter to get only columns with loaded LoRAs
    loaded_loras = []
    for idx, item in enumerate(selected_loras_state):
        if item[0] is not None:
             # item structure: (image_index, version_index, scale)
             if len(item) == 3:
                 loaded_loras.append((idx, item[0], item[1], item[2]))
             else:
                 # Legacy state support
                 loaded_loras.append((idx, item[0], None, item[1]))

    print(f"Loaded LoRAs: {loaded_loras}")
    print(f"Quantity: {quantity}")
    print(f"Selected LoRAs: {selected_loras_state}")    

    
    if not loaded_loras:
        raise gr.Error("You must select at least one LoRA before proceeding.")

    prompts = [
        p.strip() for p in [prompt_1, prompt_2, prompt_3, prompt_4]
        if p and p.strip()
    ]

    if not prompts:
        raise gr.Error("You must fill at least one prompt.")

    # limpiar LoRAs previas
    pipe.unload_lora_weights()

    # 🔥 CARGA DE MÚLTIPLES LORAs
    adapter_names = []
    adapter_weights = []
    
    # Add lightning LoRA if in speed mode
    if speed_mode == "light 4":
        pipe.load_lora_weights(
            LIGHTNING_LORA_REPO,
            weight_name=LIGHTNING_LORA_WEIGHT,
            adapter_name="lightning"
        )
        adapter_names.append("lightning")
        adapter_weights.append(1.0)
    elif speed_mode == "light 8":
        pipe.load_lora_weights(
            LIGHTNING_LORA_REPO,
            weight_name=LIGHTNING8_LORA_WEIGHT,
            adapter_name="lightning"
        )
        adapter_names.append("lightning")
        adapter_weights.append(1.0)
    elif speed_mode == "light 4 fp8":
        pipe.load_lora_weights(
            LIGHTNING_LORA_REPO,
            weight_name=LIGHTNING_FP8_4STEPS_LORA_WEIGHT,
            adapter_name="lightning"
        )
        adapter_names.append("lightning")
        adapter_weights.append(1.0)

    elif speed_mode == "Wuli-art":
        pipe.load_lora_weights(
            WULI_LORA_REPO,
            weight_name=WULI_LORA_WEIGHT,
            adapter_name="lightning"
        )
        adapter_names.append("lightning")
        adapter_weights.append(1.0)

    elif speed_mode == "Wuli2Step-art":
        pipe.load_lora_weights(
            WULI2STP_LORA_REPO,
            weight_name=WULI2STP_LORA_WEIGHT,
            adapter_name="lightning"
        )
        adapter_names.append("lightning")
        adapter_weights.append(1.0)
    
    # Load all selected LoRAs from columns
    for col_idx, image_idx, version_idx, scale in loaded_loras:
        selected_lora = loras[image_idx]
        
        # Determine weights and repo based on version
        if version_idx is not None and "versions" in selected_lora:
            version_data = selected_lora["versions"][version_idx]
            lora_path = version_data.get("repo") or selected_lora.get("repo")
            weight_name = version_data.get("weights", selected_lora.get("weights"))
            print(f"Using version '{version_data['title']}' for LoRA {selected_lora['title']}")
        else:
            lora_path = selected_lora.get("repo")
            weight_name = selected_lora.get("weights")

        adapter_name = f"lora_{col_idx}"
        
        pipe.load_lora_weights(
            lora_path,
            weight_name=weight_name,
            adapter_name=adapter_name
        )
        print(f"Loaded LoRA: {lora_path} as {adapter_name} with scale {scale}")
        adapter_names.append(adapter_name)
        adapter_weights.append(scale)
    
    # Set all adapters
    print(f"Setting adapters: {adapter_names} with weights {adapter_weights}")
    pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)

    # Colectar trigger words de todos los LoRAs cargados
    all_trigger_words = []
    for _, image_idx, version_idx, _ in loaded_loras:
        # PENDING: Could also fetch trigger words from version specific data if available
        trigger = loras[image_idx].get("trigger_word", "")
        if trigger:
            all_trigger_words.append(trigger)
    
    combined_trigger = " ".join(all_trigger_words)
    
    # Aplicar trigger words a los prompts
    final_prompts = []
    for p in prompts:
        if combined_trigger:
            final_prompts.append(f"{combined_trigger} {p}")
        else:
            final_prompts.append(p)
    
    prompts = final_prompts

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    multiplier = float(quality_multiplier.replace("x", ""))
    width = int(width * multiplier)
    height = int(height * multiplier)

    # ✅ FIX: quantity viene como index 0..3 (por type="index"), convertimos a 1..4
    real_quantity = int(quantity) + 1


    try:
        if (width * height > 1048 * 1048):
            pipe.vae.enable_tiling()
            print("VAE tiling enabled")
        else:
            pipe.vae.disable_tiling()
            print("VAE tiling disabled")

    except Exception as e:
        print(f"Could not enable VAE tiling: {e}")  


    if (history is None):
        history = []
    gallery_images = []
        
    for prompt in prompts:
        current_seed = seed
        if prompt_enhance:
            prompt = rewrite(prompt)
    
        # ✅ FIX: quantity ya no es el componente global; es un int real
        for _ in range(real_quantity):
            generator = torch.Generator(device="cuda").manual_seed(current_seed)
    
            result = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=steps,
                true_cfg_scale=cfg_scale,
                width=width,
                height=height,
                num_images_per_prompt=1,
                generator=generator,
            )
            
            img = result.images[0]
            imgtuple = (img, str(current_seed))

            gallery_images.append(imgtuple)
            
            # history persistente (acumula)
            history = [(img, str(current_seed))] + history
            history = history[:24]

            yield gallery_images, history, history, seed
            
            current_seed += 100  # separación segura
    
    #return images
    




# ... (El resto de las funciones como get_huggingface_safetensors, check_custom_model, etc., permanecen sin cambios) ...
def get_huggingface_safetensors(link):
    split_link = link.split("/")
    if len(split_link) != 2:
        raise Exception("Invalid Hugging Face repository link format.")
    print(f"Repository attempted: {split_link}")
    model_card = ModelCard.load(link)
    base_model = model_card.data.get("base_model")
    print(f"Base model: {base_model}")
    acceptable_models = {
        "Qwen/Qwen-Image",
        "Qwen/Qwen-Image-2512",
    }
    models_to_check = base_model if isinstance(base_model, list) else [base_model]
    if not any(model in acceptable_models for model in models_to_check):
        raise Exception("Not a Qwen-Image LoRA!")
    image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
    trigger_word = model_card.data.get("instance_prompt", "")
    image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
    fs = HfFileSystem()
    try:
        list_of_files = fs.ls(link, detail=False)
        safetensors_name = None
        for file in list_of_files:
            filename = file.split("/")[-1]
            if filename.endswith(".safetensors"):
                safetensors_name = filename
                break
        if not safetensors_name:
            raise Exception("No valid *.safetensors file found in the repository.")
    except Exception as e:
        print(e)
        raise Exception("You didn't include a valid Hugging Face repository with a *.safetensors LoRA")
    return split_link[1], link, safetensors_name, trigger_word, image_url


def check_custom_model(link):
    print(f"Checking a custom model on: {link}")
    if link.endswith('.safetensors'):
        if 'huggingface.co' in link:
            parts = link.split('/')
            try:
                hf_index = parts.index('huggingface.co')
                username = parts[hf_index + 1]
                repo_name = parts[hf_index + 2]
                repo = f"{username}/{repo_name}"
                safetensors_name = parts[-1]
                try:
                    model_card = ModelCard.load(repo)
                    trigger_word = model_card.data.get("instance_prompt", "")
                    image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
                    image_url = f"https://huggingface.co/{repo}/resolve/main/{image_path}" if image_path else None
                except:
                    trigger_word = ""
                    image_url = None
                return repo_name, repo, safetensors_name, trigger_word, image_url
            except:
                raise Exception("Invalid safetensors URL format")
    if link.startswith("https://"):
        if link.startswith("https://huggingface.co") or link.startswith("https://www.huggingface.co"):
            link_split = link.split("huggingface.co/")
            return get_huggingface_safetensors(link_split[1])
    else: 
        return get_huggingface_safetensors(link)


def add_custom_lora(custom_lora, *args):
    # args: radios (N) + scales (N) + selected_loras (1)
    num_radios = NUM_LORAS
    radio_vals = args[:num_radios]
    scale_vals = args[num_radios:num_radios*2]
    selected_loras_state = args[-1]
    
    global loras
    if custom_lora:
        try:
            title, repo, path, trigger_word, image = check_custom_model(custom_lora)
            print(f"Loaded custom LoRA: {repo}")
            model_card_examples = ""
            try:
                model_card = ModelCard.load(repo)
                widget_data = model_card.data.get("widget", [])
                if widget_data and len(widget_data) > 0:
                    examples_html = '<div style="margin-top: 10px;">'
                    examples_html += '<h4 style="margin-bottom: 8px; font-size: 0.9em;">Sample Images:</h4>'
                    examples_html += '<div style="display: grid; grid-template-columns: repeat(4, 1fr); gap: 8px;">'
                    for i, example in enumerate(widget_data[:4]):
                        if "output" in example and "url" in example["output"]:
                            image_url = f"https://huggingface.co/{repo}/resolve/main/{example['output']['url']}"
                            caption = example.get("text", f"Example {i+1}")
                            examples_html += f'''
                            <div style="text-align: center;">
                                <img src="{image_url}" style="width: 100%; height: auto; border-radius: 4px;" />
                                <p style="font-size: 0.7em; margin: 2px 0;">{caption[:30]}{'...' if len(caption) > 30 else ''}</p>
                            </div>
                            '''
                    examples_html += '</div></div>'
                    model_card_examples = examples_html
            except Exception as e:
                print(f"Could not load model card examples for custom LoRA: {e}")
            card = f'''
            <div class="custom_lora_card">
              <span>Loaded custom LoRA:</span>
              <div class="card_internal">
                <img src="{image}" />
                <div>
                    <h3>{title}</h3>
                    <small>{"Using: <code><b>"+trigger_word+"</code></b> as the trigger word" if trigger_word else "No trigger word found. If there's a trigger word, include it in your prompt"}<br></small>
                </div>
              </div>
              {model_card_examples}
            </div>
            '''
            existing_item_index = next((index for (index, item) in enumerate(loras) if item['repo'] == repo), None)
            if existing_item_index is None:
                new_item = {"image": image, "title": title, "repo": repo, "weights": path, "trigger_word": trigger_word}
                print(new_item)
                loras.append(new_item)
                existing_item_index = len(loras) - 1
            
            # Update the selected slot with the custom LoRA
            md_updates, new_state, scale_updates, del_btn_updates, dropdown_updates, _ = update_slot_with_lora(
                existing_item_index, radio_vals, scale_vals, selected_loras_state
            )
            
            return (
                gr.update(visible=True, value=card), 
                gr.update(visible=True), 
                reload_loras_gallery()
            ) + tuple(md_updates) + (new_state,) + tuple(scale_updates) + tuple(del_btn_updates) + tuple(dropdown_updates)
            
        except Exception as e:
            full_traceback = traceback.format_exc()
            print(f"Full traceback:\n{full_traceback}")
            gr.Warning(f"Invalid LoRA: either you entered an invalid link, or a non-Qwen-Image LoRA, this was the issue: {e}")
            
            empty_updates = [gr.update()] * num_radios
            return (
                gr.update(visible=True, value=f"Invalid LoRA: either you entered an invalid link, a non-Qwen-Image LoRA"), 
                gr.update(visible=True), 
                gr.update()
            ) + tuple(empty_updates) + (selected_loras_state,) + tuple(empty_updates) + tuple(empty_updates) + tuple(empty_updates)
    else:
        empty_updates = [gr.update()] * num_radios
        return (
            gr.update(visible=False), 
            gr.update(visible=False), 
            gr.update()
        ) + tuple(empty_updates) + (selected_loras_state,) + tuple(empty_updates) + tuple(empty_updates) + tuple(empty_updates)


def remove_custom_lora():
    return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""


def reload_loras_gallery():
    global loras
    loras = load_loras_hot()

    gallery_items = [
        (item["image"], item.get("title") or item.get("name"))
        for item in loras
        if item.get("image")
    ]

    return gr.update(value=gallery_items)



def init(speed_mode, aspect_ratio):
    loras_result = reload_loras_gallery()
    speed_mode_result = handle_speed_mode(speed_mode)
    aspect_ratio_result = apply_aspect_ratio(aspect_ratio)
    return (
        *speed_mode_result,
        *aspect_ratio_result,
        loras_result,
        "1"  # Preselect first radio
    )



def update_slider_state(val, state, idx):
    # state is list of tuples (image_index, version_index, scale)
    new_state = [tuple(item) for item in state]
    
    if len(new_state[idx]) == 3:
        current_image, current_version, current_scale = new_state[idx]
    else:
        current_image, current_scale = new_state[idx]
        current_version = None
        
    new_state[idx] = (current_image, current_version, float(val))
    print(f"Slider {idx} updated. New state: {new_state}") 
    return new_state

def remove_lora(state, idx):
    new_state = list(state)
    new_state[idx] = (None, None, 1.0) # Reset to default
    
    # Return updates for: global_state, radio, slider, markdown, delete_btn, output_text, version_dropdown
    return (
        new_state,                 # selected_loras
        gr.update(value=None),     # radio
        gr.update(visible=False, value=1.0), # slider
        gr.update(value=""),       # markdown
        gr.update(visible=False),  # delete_btn
        gr.update(visible=False, value=None, choices=[]) # version_dropdown
    )

def validate_generate_button(state):
    """Enable generate button if at least one LoRA is loaded."""
    # state is list of tuples (image_index, version_index, scale)
    # Check if any column has a LoRA loaded (image_index is not None)
    has_lora = any(item[0] is not None for item in state)
    return gr.update(interactive=has_lora)

def update_slot_with_lora(lora_index, radio_vals, scale_vals, current_state):
    """
    Helper function to update a specific slot with a selected LoRA index.
    """
    num_radios = NUM_LORAS
    
    # Identify index of selected radio
    selected_slot_index = -1
    for i, val in enumerate(radio_vals):
        if val is not None:
            selected_slot_index = i
            break
            
    # Prepare update arrays
    md_updates = [gr.update() for _ in range(num_radios)]
    scale_updates = [gr.update() for _ in range(num_radios)]
    del_btn_updates = [gr.update() for _ in range(num_radios)]
    dropdown_updates = [gr.update() for _ in range(num_radios)]
    
    new_state = list(current_state)
    
    if selected_slot_index != -1 and lora_index is not None:
        # Get LoRA details
        lora = loras[lora_index]
        lora_name = lora["title"]
        lora_image = lora["image"]
        
        versions = lora.get("versions", [])
        
        version_index = None
        version_title = ""

        if versions:
            # Has versions. Find best version or default to last.
            found_best = False
            version_index = -1
            
            for i, v in enumerate(versions, start=0):
                if v.get("best", False):
                    version_index = i
                    found_best = True
                    break
            
            if not found_best:
                 # Default to the last one if not found
                 version_index = len(versions) - 1
            
            version_data = versions[version_index]
            version_title = version_data["title"]
            version_image = version_data.get("image", lora_image)
            
            version_choices = [v["title"] for v in versions]
            
            dropdown_updates[selected_slot_index] = gr.update(visible=True, choices=version_choices, value=version_title)
        else:
            dropdown_updates[selected_slot_index] = gr.update(visible=False, value=None, choices=[])


        # Update specific markdown with image and name
        markdown_content = f"<img src='{version_image if versions else lora_image}' style='height:100px; display:block; margin-bottom:10px;' />\n\n**{lora_name}**{' ('+version_title+')' if version_title else ''}"
        md_updates[selected_slot_index] = gr.update(value=markdown_content)
        scale_updates[selected_slot_index] = gr.update(visible=True)
        del_btn_updates[selected_slot_index] = gr.update(visible=True)
        
        # Update state for this column: (image_index, version_index, scale)
        if len(scale_vals) > selected_slot_index:
             current_scale = scale_vals[selected_slot_index]
        else:
             current_scale = 1.0 # Fallback
             
        new_state[selected_slot_index] = (lora_index, version_index, current_scale)

    return md_updates, new_state, scale_updates, del_btn_updates, dropdown_updates, selected_slot_index


def get_selection(evt: gr.SelectData, *args):
    # args: radios (N) + scales (N) + mds (N) + dropdowns (N) + selected_loras (1) + gallery (1)
    
    num_radios = NUM_LORAS
    radio_vals = args[:num_radios]
    scale_vals = args[num_radios:num_radios*2]
    # mds skipped in logic
    # dropdowns skipped in logic
    current_state = args[-2]
    # gallery_val skipped
    
    md_updates, new_state, scale_updates, del_btn_updates, dropdown_updates, _ = update_slot_with_lora(
        evt.index, radio_vals, scale_vals, current_state
    )

    return md_updates + [new_state] + scale_updates + del_btn_updates + dropdown_updates




css = '''
#gen_btn{height: 100%}
#gen_column{align-self: stretch}
#title{text-align: center}
#title h1{font-size: 3em; display:inline-flex; align-items:center}
#title img{width: 100px; margin-right: 0.5em}
#gallery .grid-wrap{height: 10vh}
#lora_list{background: var(--block-background-fill);padding: 0 1em .3em; font-size: 90%}
.card_internal{display: flex;height: 100px;margin-top: .5em}
.card_internal img{margin-right: 1em}
.styler{--form-gap-width: 0px !important}
#speed_status{padding: .5em; border-radius: 5px; margin: 1em 0}
'''


with gr.Blocks(theme=gr.themes.Soft(), css=css, delete_cache=(60, 60)) as app:
    title = gr.HTML(
        """<h1 style=\"color:#644fea\">Qwen-Image-2512</h1>        
        <h3 style=\"margin-top: -10px\">LoRA🦜 ChoquinLabs Explorer</h3>""",
        elem_id="title",
    )

    # Login opcional: si el visitante entra con su cuenta, el upscale 4x se
    # cobra a su cupo de ZeroGPU en vez de salir anónimo (cupo por IP).
    gr.LoginButton(size="sm")

    selected_index = gr.State(None)
    
    with gr.Row():
        with gr.Column(scale=3):
    
            prompt_1 = gr.Textbox(label="Prompt 1", lines=1)
            prompt_2 = gr.Textbox(label="Prompt 2", lines=1)
            prompt_3 = gr.Textbox(label="Prompt 3", lines=1)
            prompt_4 = gr.Textbox(label="Prompt 4", lines=1)
            
            negative_prompt = gr.Textbox(label="Negative Prompt", lines=1, placeholder="Optional: what to avoid")
            prompt_enhance = gr.Checkbox(label="Prompt Enhance", value=False)
        with gr.Column(scale=1, elem_id="gen_column"):
            generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")


    

    selected_loras = gr.State([(None, None, 1.0)] * NUM_LORAS)
    radios = []
    scales = []
    mds = []
    delete_btns = []
    version_dropdowns = []


    with gr.Row():
        with gr.Column():


            with gr.Row():
                for i in range(NUM_LORAS):
                    with gr.Column():
                        # Each radio has a single choice which is its column number
                        r = gr.Radio(
                            [str(i + 1)], 
                            label=f"Lora {i + 1}"
                        )
                        radios.append(r)
                        
                        md = gr.Markdown("")
                        mds.append(md)
                        
                        # New version dropdown
                        v_dd = gr.Dropdown(label="Version", choices=[], visible=False, interactive=True)
                        version_dropdowns.append(v_dd)

                        lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=3, step=0.1, value=1.0, interactive=True, visible=False)
                        scales.append(lora_scale)
                        
                        del_btn = gr.Button("🗑️", visible=False)
                        delete_btns.append(del_btn)



            
            selected_info = gr.Markdown("")
            examples_component = gr.Examples(examples=[], inputs=[prompt_1], label="Sample Prompts", visible=False)
            gallery = gr.Gallery(
                [(item["image"], item["title"]) for item in loras],
                label="LoRA Gallery",
                allow_preview=False,
                columns=3,
                elem_id="gallery"
            )
            reload_btn = gr.Button("🔄 Reload LoRAs")            
            
            
            with gr.Group():
                custom_lora = gr.Textbox(label="Custom LoRA", info="LoRA Hugging Face path", placeholder="username/qwen-image-custom-lora")
                gr.Markdown("[Check Qwen-Image LoRAs](https://huggingface.co/models?other=base_model:adapter:Qwen/Qwen-Image)", elem_id="lora_list")
            custom_lora_info = gr.HTML(visible=False)
            custom_lora_button = gr.Button("Remove custom LoRA", visible=False)
        
        with gr.Column():
            result = gr.Gallery(label="Generated Images", show_label=True, elem_id="result_gallery")
            history_state = gr.State([]) 
            ### MODIFICACIÓN 2: AÑADIR LOS COMPONENTES DE LA UI DEL HISTORIAL ###
            with gr.Group():
                with gr.Row():
                    gr.Markdown("### 📜 History")
                    clear_history_button = gr.Button("🗑️ Clear History", size="sm")


                history_gallery = gr.Gallery(
                    label="Generation History",
                    show_label=False,
                    columns=4,
                    object_fit="contain",
                    height="auto",
                    interactive=False
                )

                selected_history_image = gr.State(None)
                upscale_button = gr.Button("🔍 Upscale 4x (PiD)", size="sm")
                upscaled_slider = gr.ImageSlider(
                    label="Original ↔ PiD 4x upscale",
                    visible=False
                )
            ### FIN DE LA MODIFICACIÓN 2 ###
            
            with gr.Row():
                with gr.Column():
                    speed_mode = gr.Radio(
                        label="Generation Mode",
                        choices=["light 4", "Wuli-art", "Wuli2Step-art", "light 4 fp8", "light 8", "normal"],
                        value="light 4",
                        info="'light' modes use Lightning LoRA for faster generation"
                    )
                with gr.Column():
                    quantity = gr.Radio(
                        label="Quantity",
                        choices=["1", "2", "3", "4"],
                        value="1",
                        type="index"
                    )
            
            speed_status = gr.Markdown("Quality mode active", elem_id="speed_status")

            with gr.Row():
                aspect_ratio = gr.Radio(
                    label="Aspect Ratio",
                    choices=["1:1", "16:9", "9:16", "3:2", "2:3", "3:1", "2:1"],
                    value="16:9"
                )


            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=1920,
                    step=1,
                    value=1920
                )
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=1920,
                    step=1,
                    value=1080
                )


            with gr.Row():
                quality_multiplier = gr.Radio(
                    label="Quality (Size Multiplier)",
                    choices=["0.5x", "0.75x", "1x", "1.25x", "1.5x", "2x"],
                    value="1x"
                )


    with gr.Row():
        with gr.Accordion("Advanced Settings", open=False):
            with gr.Column():
                with gr.Row():
                    cfg_scale = gr.Slider(
                        label="Guidance Scale (True CFG)", 
                        minimum=1.0, 
                        maximum=5.0, 
                        step=0.1, 
                        value=3.5,
                        info="Lower for speed mode, higher for quality"
                    )
                    steps = gr.Slider(
                        label="Steps", 
                        minimum=1, 
                        maximum=50, 
                        step=1, 
                        value=45,
                        info="Automatically set by speed mode"
                    )
                
                with gr.Row():
                    randomize_seed = gr.Checkbox(True, label="Randomize seed")
                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)


    # Event handlers
    # gallery.select(
    #     update_selection,
    #     inputs=[width, height],
    #     outputs=[prompt_1, selected_info, selected_index, width, height, generate_button]
    # )
    
    speed_mode.change(
        handle_speed_mode,
        inputs=[speed_mode],
        outputs=[speed_status, steps, cfg_scale]
    )
    
    custom_lora.submit(
        fn=add_custom_lora,
        inputs=[custom_lora] + radios + scales + [selected_loras],
        outputs=[custom_lora_info, custom_lora_button, gallery] + mds + [selected_loras] + scales + delete_btns + version_dropdowns
    )
    
    custom_lora_button.click(
        remove_custom_lora,
        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, custom_lora]
    )
    
    ### MODIFICACIÓN 3: CONECTAR LOS EVENTOS DEL HISTORIAL ###
    # Evento principal de generación
    generate_event = gr.on(
        triggers=[generate_button.click, prompt_1.submit],
    
        fn=run_lora_multi,
        inputs=[
            prompt_1, prompt_2, prompt_3, prompt_4,
            negative_prompt,
            cfg_scale, steps, selected_loras,  # Changed from selected_index
            randomize_seed, seed,
            width, height,  # Removed lora_scale
            speed_mode, quality_multiplier,
            quantity,
            history_state,
            prompt_enhance
        ],
        outputs=[result, history_gallery, history_state, seed]
    )


    
    # Encadenar la actualización del historial para que se ejecute DESPUÉS de la generación



    # Evento para el botón de limpiar historial
    clear_history_button.click(
        fn=clear_history,
        inputs=None,
        outputs=[history_state, history_gallery]
    )

    # Upscale 4x de cualquier imagen del historial, vía API del space PiD
    history_gallery.select(
        fn=select_history_image,
        inputs=[history_state],
        outputs=[selected_history_image]
    )

    upscale_button.click(
        fn=upscale_history_image,
        inputs=[selected_history_image],
        outputs=[upscaled_slider]
    )
    ### FIN DE LA MODIFICACIÓN 3 ###

    aspect_ratio.change(
        fn=apply_aspect_ratio,
        inputs=[aspect_ratio],
        outputs=[width, height]
    )

    reload_btn.click(
        fn=reload_loras_gallery,
        outputs=gallery,
    )



    for i, r in enumerate(radios):
        others = radios[:i] + radios[i+1:]
        # JS: if val is selected (true), return nulls for all others. 
        # Otherwise return current values (no change).
        js_code = f"(val, ...args) => val ? args.map(_ => null) : args"
        r.change(fn=None, inputs=[r] + others, outputs=others, js=js_code)
        
        # JS toggle for gallery class
        js_gallery_toggle = "(...args) => { const gallery = document.getElementById('gallery'); const anySelected = args.some(v => v !== null && v !== ''); if (gallery) { if (anySelected) gallery.classList.remove('disabled'); else gallery.classList.add('disabled'); } }"
        r.change(fn=None, inputs=radios, outputs=None, js=js_gallery_toggle)
        
    # Bind slider changes separately to ensure all inputs/outputs are available
    for i in range(NUM_LORAS):
        # Slider change
        scales[i].change(
            fn=partial(update_slider_state, idx=i),
            inputs=[scales[i], selected_loras],
            outputs=[selected_loras]
        )
        scales[i].change(fn=validate_generate_button, inputs=[selected_loras], outputs=[generate_button])
        
        # Delete LoRA
        delete_btns[i].click(
            fn=partial(remove_lora, idx=i),
            inputs=[selected_loras],
            outputs=[selected_loras, radios[i], scales[i], mds[i], delete_btns[i], version_dropdowns[i]]
        )
        delete_btns[i].click(fn=validate_generate_button, inputs=[selected_loras], outputs=[generate_button])

        # Version change
        version_dropdowns[i].change(
            fn=partial(update_version_selection, idx=i),
            inputs=[version_dropdowns[i], selected_loras],
            outputs=[mds[i], selected_loras]
        )

    
    # When gallery is selected
    gallery.select(
        fn=get_selection,
        inputs=radios + scales + mds + version_dropdowns + [selected_loras, gallery], # Pass dropdowns too
        outputs=mds + [selected_loras] + scales + delete_btns + version_dropdowns
    )
    
    # Also trigger validation on selection
    gallery.select(fn=validate_generate_button, inputs=[selected_loras], outputs=[generate_button])

    
    # -----------------------------------------------------
    # START GENERATION
    # -----------------------------------------------------
    # Handler is already defined in generate_event (gr.on triggers)



    load_event = app.load(
        fn=init,
        inputs=[gr.State("light 4"), gr.State(DEFAULT_ASPECT_RATIO)],
        outputs=[speed_status, steps, cfg_scale, width, height, gallery, radios[0]]
    )
    load_event.then(fn=validate_generate_button, inputs=[selected_loras], outputs=[generate_button])



app.queue()
app.launch()