import spaces
import gradio as gr
import numpy as np
import PIL.Image
from PIL import Image
import random
from diffusers import StableDiffusionXLPipeline
from diffusers import EulerAncestralDiscreteScheduler
import torch
from compel import Compel, ReturnedEmbeddingsType
from huggingface_hub import login, HfApi
import os

# Add your Hugging Face token here or set it as an environment variable
HF_TOKEN = os.getenv("HF_TOKEN")

if HF_TOKEN:
    login(token=HF_TOKEN)

# --- LoRA Mapping ---
LORA_MAPPING = {
    "LoCon_d128-128_a16-32_n151_b4_lr=3e-04-5e-05_joycaption_seed=100-20": {
        "repo": "rfyuan/waiREALCN_v14_LoRA",
        "file": "LoCon_d128.128_a16.32_n151_b4-lr=3.00e-04-5.00e-05_joycaption_seed=100-20.safetensors"
    },
    "LoCon_d128-128_a16-32_n151_b4_lr=5e-04-5e-05_joycaption_seed=100-18": {
        "repo": "rfyuan/waiREALCN_v14_LoRA",
        "file": "LoCon_d128.128_a16.32_n151_b4-lr=5.00e-04-5.00e-05_joycaption_seed=100-18.safetensors"
    },
}
# --- End LoRA Mapping ---

# --- Define a single repository for all dynamic LoRAs ---
DYNAMIC_LORA_REPO = "rfyuan/waiREALCN_v14_LoRA"


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
pipe = None
compel = None
model_loaded = False
FAILED_LORAS = set()
AVAILABLE_DYNAMIC_LORAS = []

try:
    pipe = StableDiffusionXLPipeline.from_pretrained(
        "rfyuan/waiREALCN_v14_usdf",
        torch_dtype=torch.float16,
        variant="fp16",
        use_safetensors=True,
    )
    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
    pipe.to(device)

    compel = Compel(
        tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
        text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
        returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
        requires_pooled=[False, True],
        truncate_long_prompts=False
    )
    
    model_loaded = True
except Exception as e:
    print(f"Failed to load model: {e}")

# --- Fetch dynamic LoRAs from the specified repo at startup ---
if model_loaded:
    print(f"Fetching available LoRAs from {DYNAMIC_LORA_REPO}...")
    try:
        api = HfApi()
        repo_files = api.list_repo_files(repo_id=DYNAMIC_LORA_REPO, repo_type="model")
        AVAILABLE_DYNAMIC_LORAS = [f for f in repo_files if f.endswith(".safetensors")]
        print(f"Found {len(AVAILABLE_DYNAMIC_LORAS)} available LoRAs.")
    except Exception as e:
        print(f"Failed to fetch LoRAs from repo: {e}")

# --- PRE-DOWNLOADING ONLY FIXED LORAS AT STARTUP ---
if model_loaded:
    print("Pre-downloading fixed LoRAs...")
    for name, data in LORA_MAPPING.items():
        try:
            pipe.load_lora_weights(data["repo"], weight_name=data["file"], adapter_name=name)
            print(f"Successfully cached LoRA: {name}")
        except Exception as e:
            print(f"Failed to cache LoRA '{name}': {e}")
            FAILED_LORAS.add(name)
    
    print("Unloading all LoRAs from VRAM after caching.")
    pipe.unload_lora_weights()


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

def process_long_prompt(prompt, negative_prompt=""):
    try:
        conditioning, pooled = compel([prompt, negative_prompt])
        return conditioning, pooled
    except Exception as e:
        print(f"Long prompt processing failed: {e}, falling back to standard processing")
        return None, None

# --- NEW FUNCTION TO REFRESH THE LORA LIST ---
def refresh_lora_list():
    print("Refreshing dynamic LoRA list...")
    try:
        api = HfApi()
        repo_files = api.list_repo_files(repo_id=DYNAMIC_LORA_REPO, repo_type="model")
        
        global AVAILABLE_DYNAMIC_LORAS
        AVAILABLE_DYNAMIC_LORAS = [f for f in repo_files if f.endswith(".safetensors")]
        
        print(f"Found {len(AVAILABLE_DYNAMIC_LORAS)} available LoRAs.")
        return gr.update(choices=["None"] + AVAILABLE_DYNAMIC_LORAS)
    except Exception as e:
        print(f"Failed to refresh LoRAs from repo: {e}")
        return gr.update() # Return an empty update to not change the UI on error

def select_dynamic_lora(lora_name):
    if not lora_name or lora_name == "None":
        return None, gr.update(visible=False), "No dynamic LoRA selected."

    adapter_name = "dynamic_lora_cache_check"
    
    try:
        print(f"Pre-caching dynamic LoRA: {lora_name}")
        pipe.load_lora_weights(DYNAMIC_LORA_REPO, weight_name=lora_name, adapter_name=adapter_name)
        pipe.unload_lora_weights()
        
        status_message = f"✅ LoRA '{lora_name}' is ready to use."
        return lora_name, gr.update(label=lora_name, value=0.8, visible=True), status_message
        
    except Exception as e:
        print(f"Failed to pre-cache dynamic LoRA {lora_name}: {e}")
        status_message = f"Error: Could not cache LoRA '{lora_name}'."
        return None, gr.update(visible=False), status_message


@spaces.GPU(duration=30)
def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, dynamic_lora_name, dynamic_lora_weight, *lora_weights):
    if not model_loaded:
        error_img = Image.new('RGB', (width, height), color=(50, 50, 50))
        return error_img

    pipe.unload_lora_weights()
    pipe.disable_lora()

    active_loras = []
    active_weights = []

    # 1. Load pre-defined LoRAs from sliders
    for i, lora_name in enumerate(LORA_MAPPING.keys()):
        if lora_name in FAILED_LORAS:
            continue
        weight = lora_weights[i]
        if weight > 0:
            try:
                data = LORA_MAPPING[lora_name]
                print(f"Loading pre-defined LoRA: {lora_name}")
                pipe.load_lora_weights(data["repo"], weight_name=data["file"], adapter_name=lora_name)
                active_loras.append(lora_name)
                active_weights.append(weight)
            except Exception as e:
                print(f"Failed to load LoRA {lora_name} from cache: {e}")
                continue
    
    # Load the dynamic LoRA if selected
    if dynamic_lora_name and dynamic_lora_name != "None" and dynamic_lora_weight > 0:
        try:
            adapter_name = "dynamic_lora"
            print(f"Loading dynamic LoRA from {DYNAMIC_LORA_REPO}: {dynamic_lora_name}")
            pipe.load_lora_weights(DYNAMIC_LORA_REPO, weight_name=dynamic_lora_name, adapter_name=adapter_name)
            active_loras.append(adapter_name)
            active_weights.append(dynamic_lora_weight)
        except Exception as e:
            print(f"Failed to load dynamic LoRA {dynamic_lora_name} from cache: {e}")


    try:
        # 2. Set the weights for all active adapters.
        if active_loras:
            print(f"Activating LoRAs: {list(zip(active_loras, active_weights))}")
            pipe.set_adapters(active_loras, adapter_weights=active_weights)

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

        generator = torch.Generator(device=device).manual_seed(seed)
        
        # 3. Generate the image
        use_long_prompt = len(prompt.split()) > 10 or len(prompt) > 200
        if use_long_prompt:
            conditioning, pooled = process_long_prompt(prompt, negative_prompt)
            if conditioning is not None:
                output_image = pipe(
                    prompt_embeds=conditioning[0:1],
                    pooled_prompt_embeds=pooled[0:1],
                    negative_prompt_embeds=conditioning[1:2],
                    negative_pooled_prompt_embeds=pooled[1:2],
                    guidance_scale=guidance_scale,
                    num_inference_steps=num_inference_steps,
                    width=width,
                    height=height,
                    generator=generator,
                ).images[0]
                return output_image
        
        output_image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator
        ).images[0]
        
        return output_image
    except Exception as e:
        print(f"Error during generation: {e}")
        error_img = Image.new('RGB', (width, height), color=(0, 0, 0))
        return error_img
    finally:
        # 4. Unload all LoRAs to free up VRAM for the next user.
        print("Unloading LoRAs to free VRAM.")
        pipe.unload_lora_weights()
        pipe.disable_lora()


css = """
#col-container {
    margin: 0 auto;
    max-width: 768px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        if not model_loaded:
            gr.Markdown("⚠️ **Model failed to load. Please check logs for errors.**")

        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                lines=3,
                placeholder="Enter your prompt",
                container=False,
            )
            run_button = gr.Button("Run", scale=0)

        result = gr.Image(label="Result", show_label=False)
        
        with gr.Accordion("Advanced Settings", open=False):
            with gr.Group():
                gr.Markdown("### Select Dynamic LoRA")
                with gr.Row():
                    dynamic_lora_dropdown = gr.Dropdown(
                        choices=["None"] + AVAILABLE_DYNAMIC_LORAS,
                        value="None",
                        label="Available Dynamic LoRAs",
                        scale=4
                    )
                    # --- NEW: Refresh button ---
                    refresh_button = gr.Button("Refresh", scale=1)

                dynamic_lora_status = gr.Markdown()
                dynamic_lora_state = gr.State(None) 
            
            with gr.Group():
                gr.Markdown("### LoRA Weights (0 = Off)")
                lora_sliders = []
                for name in LORA_MAPPING.keys():
                    if name in FAILED_LORAS:
                        continue
                    slider = gr.Slider(
                        label=name,
                        minimum=0.0,
                        maximum=2.0,
                        step=0.05,
                        value=0.0
                    )
                    lora_sliders.append(slider)
                
                dynamic_lora_slider = gr.Slider(
                    label="Dynamic LoRA",
                    minimum=0.0,
                    maximum=2.0,
                    step=0.05,
                    value=0.8,
                    visible=False
                )

            negative_prompt = gr.Text(
                label="Negative prompt",
                max_lines=1,
                placeholder="Enter a negative prompt",
                value="(low quality, worst quality)1.2, very displeasing, 3d, watermark, signature, ugly, poorly drawn"
            )

            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )

            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.0,
                    maximum=20.0,
                    step=0.1,
                    value=7,
                )
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=28,
                    step=1,
                    value=28,
                )
    
    # --- MODIFIED: Wire up the dynamic LoRA dropdown and refresh button ---
    dynamic_lora_dropdown.change(
        fn=select_dynamic_lora,
        inputs=[dynamic_lora_dropdown],
        outputs=[dynamic_lora_state, dynamic_lora_slider, dynamic_lora_status]
    )
    
    refresh_button.click(
        fn=refresh_lora_list,
        inputs=None,
        outputs=[dynamic_lora_dropdown]
    )

    run_button.click(
        fn=infer,
        inputs=[
            prompt, negative_prompt, seed, randomize_seed, 
            width, height, guidance_scale, num_inference_steps,
            dynamic_lora_state, dynamic_lora_slider
        ] + lora_sliders,
        outputs=[result]
    )

demo.queue().launch()