import gradio as gr
import numpy as np
import spaces
import torch
import random
import json
import os
from PIL import Image
from diffusers import FluxKontextPipeline
from diffusers.utils import load_image
from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard
from safetensors.torch import load_file
import requests
import re

# Load Kontext model
MAX_SEED = np.iinfo(np.int32).max

pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")

# Load LoRA data (you'll need to create this JSON file or modify to load your LoRAs)

with open("flux_loras.json", "r") as file:
    data = json.load(file)
    flux_loras_raw = [
        {
            "image": item["image"],
            "title": item["title"],
            "repo": item["repo"],
            "trigger_word": item.get("trigger_word", ""),
            "trigger_position": item.get("trigger_position", "prepend"),
            "weights": item.get("weights", "pytorch_lora_weights.safetensors"),
        }
        for item in data
    ]
print(f"Loaded {len(flux_loras_raw)} LoRAs from JSON")
# Global variables for LoRA management
current_lora = None
lora_cache = {}

def load_lora_weights(repo_id, weights_filename):
    """Load LoRA weights from HuggingFace"""
    try:
        if repo_id not in lora_cache:
            lora_path = hf_hub_download(repo_id=repo_id, filename=weights_filename)
            lora_cache[repo_id] = lora_path
        return lora_cache[repo_id]
    except Exception as e:
        print(f"Error loading LoRA from {repo_id}: {e}")
        return None

def update_selection(selected_state: gr.SelectData, flux_loras):
    """Update UI when a LoRA is selected"""
    if selected_state.index >= len(flux_loras):
        return "### No LoRA selected", gr.update(), None
    
    lora_repo = flux_loras[selected_state.index]["repo"]
    trigger_word = flux_loras[selected_state.index]["trigger_word"]
    
    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo})"
    new_placeholder = f"optional description, e.g. 'a man with glasses and a beard'"
    
    return updated_text, gr.update(placeholder=new_placeholder), selected_state.index

def get_huggingface_lora(link):
    """Download LoRA from HuggingFace link"""
    split_link = link.split("/")
    if len(split_link) == 2:
        try:
            model_card = ModelCard.load(link)
            trigger_word = model_card.data.get("instance_prompt", "")
            
            fs = HfFileSystem()
            list_of_files = fs.ls(link, detail=False)
            safetensors_file = None
            
            for file in list_of_files:
                if file.endswith(".safetensors") and "lora" in file.lower():
                    safetensors_file = file.split("/")[-1]
                    break
            
            if not safetensors_file:
                safetensors_file = "pytorch_lora_weights.safetensors"
            
            return split_link[1], safetensors_file, trigger_word
        except Exception as e:
            raise Exception(f"Error loading LoRA: {e}")
    else:
        raise Exception("Invalid HuggingFace repository format")

def load_custom_lora(link):
    """Load custom LoRA from user input"""
    if not link:
        return gr.update(visible=False), "", gr.update(visible=False), None, gr.Gallery(selected_index=None), "### Click on a LoRA in the gallery to select it", None
    
    try:
        repo_name, weights_file, trigger_word = get_huggingface_lora(link)
        
        card = f'''
        <div style="border: 1px solid #ddd; padding: 10px; border-radius: 8px; margin: 10px 0;">
            <span><strong>Loaded custom LoRA:</strong></span>
            <div style="margin-top: 8px;">
                <h4>{repo_name}</h4>
                <small>{"Using: <code><b>"+trigger_word+"</b></code> as trigger word" if trigger_word else "No trigger word found"}</small>
            </div>
        </div>
        '''
        
        custom_lora_data = {
            "repo": link,
            "weights": weights_file,
            "trigger_word": trigger_word
        }
        
        return gr.update(visible=True), card, gr.update(visible=True), custom_lora_data, gr.Gallery(selected_index=None), f"Custom: {repo_name}", None
    
    except Exception as e:
        return gr.update(visible=True), f"Error: {str(e)}", gr.update(visible=False), None, gr.update(), "### Click on a LoRA in the gallery to select it", None

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

def classify_gallery(flux_loras):
    """Sort gallery by likes"""
    sorted_gallery = sorted(flux_loras, key=lambda x: x.get("likes", 0), reverse=True)
    return [(item["image"], item["title"]) for item in sorted_gallery], sorted_gallery

def infer_with_lora_wrapper(input_image, prompt, selected_index, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.75, flux_loras=None, progress=gr.Progress(track_tqdm=True)):
    """Wrapper function to handle state serialization"""
    return infer_with_lora(input_image, prompt, selected_index, custom_lora, seed, randomize_seed, guidance_scale, lora_scale, flux_loras, progress)

@spaces.GPU
def infer_with_lora(input_image, prompt, selected_index, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.0, flux_loras=None, progress=gr.Progress(track_tqdm=True)):
    """Generate image with selected LoRA"""
    global current_lora, pipe
    
    # Check if input image is provided
    if input_image is None:
        gr.Warning("Please upload an image first!")
        return None, seed, gr.update(visible=False)
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    # Determine which LoRA to use
    lora_to_use = None
    if custom_lora:
        lora_to_use = custom_lora
    elif selected_index is not None and flux_loras and selected_index < len(flux_loras):
        lora_to_use = flux_loras[selected_index]
    print(f"Loaded {len(flux_loras)} LoRAs from JSON")
    # Load LoRA if needed
    if lora_to_use and lora_to_use != current_lora:
        try:
            # Unload current LoRA
            if current_lora:
                pipe.unload_lora_weights()
            
            # Load new LoRA
            lora_path = load_lora_weights(lora_to_use["repo"], lora_to_use["weights"])
            if lora_path:
                pipe.load_lora_weights(lora_path, adapter_name="selected_lora")
                pipe.set_adapters(["selected_lora"], adapter_weights=[lora_scale])
                print(f"loaded: {lora_path} with scale {lora_scale}")
                current_lora = lora_to_use
            
        except Exception as e:
            print(f"Error loading LoRA: {e}")
            # Continue without LoRA
    else:
        print(f"using already loaded lora: {lora_to_use}")
    
    input_image = input_image.convert("RGB")
    
    # Check if LoRA is selected
    if lora_to_use is None:
        gr.Warning("Please select a LoRA style from the gallery first!")
        return None, seed, gr.update(visible=False)
    
    # Add trigger word to prompt
    trigger_word = lora_to_use.get("trigger_word", "")
    if trigger_word == ", How2Draw":
        prompt = f"create a How2Draw sketch of the person of the photo {prompt}, maintain the facial identity of the person and general features"
    elif trigger_word == ", video game screenshot in the style of THSMS":
        prompt = f"create a video game screenshot in the style of THSMS with the person from the photo, {prompt}. maintain the facial identity of the person and general features"
    else:
        prompt = f"convert the style of this portrait photo to {trigger_word} while maintaining the identity of the person. {prompt}. Make sure to maintain the person's facial identity and features, while still changing the overall style to {trigger_word}."
    
    try:
        image = pipe(
            image=input_image, 
            prompt=prompt,
            guidance_scale=guidance_scale,
            generator=torch.Generator().manual_seed(seed),
        ).images[0]
        
        return image, seed, gr.update(visible=True)
    
    except Exception as e:
        print(f"Error during inference: {e}")
        return None, seed, gr.update(visible=False)

# CSS styling with beautiful gradient pastel design
css = """
/* Global background and container styling */
.gradio-container {
    background: linear-gradient(135deg, #ffeef8 0%, #e6f3ff 25%, #fff4e6 50%, #f0e6ff 75%, #e6fff9 100%);
    font-family: 'Inter', sans-serif;
}

/* Main app container */
#main_app {
    display: flex;
    gap: 24px;
    padding: 20px;
    background: rgba(255, 255, 255, 0.85);
    backdrop-filter: blur(20px);
    border-radius: 24px;
    box-shadow: 0 10px 40px rgba(0, 0, 0, 0.08);
}

/* Box column styling */
#box_column {
    min-width: 400px;
}

/* Gallery box with glassmorphism */
#gallery_box {
    background: linear-gradient(135deg, rgba(255, 255, 255, 0.9) 0%, rgba(240, 248, 255, 0.9) 100%);
    border-radius: 20px;
    padding: 20px;
    box-shadow: 0 8px 32px rgba(135, 206, 250, 0.2);
    border: 1px solid rgba(255, 255, 255, 0.8);
}

/* Input image styling */
.image-container {
    border-radius: 16px;
    overflow: hidden;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.1);
}

/* Gallery styling */
#gallery {
    overflow-y: scroll !important;
    max-height: 400px;
    padding: 12px;
    background: rgba(255, 255, 255, 0.5);
    border-radius: 16px;
    scrollbar-width: thin;
    scrollbar-color: #ddd6fe #f5f3ff;
}

#gallery::-webkit-scrollbar {
    width: 8px;
}

#gallery::-webkit-scrollbar-track {
    background: #f5f3ff;
    border-radius: 10px;
}

#gallery::-webkit-scrollbar-thumb {
    background: linear-gradient(180deg, #c7d2fe 0%, #ddd6fe 100%);
    border-radius: 10px;
}

/* Selected LoRA text */
#selected_lora {
    background: linear-gradient(135deg, #818cf8 0%, #a78bfa 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    font-weight: 700;
    font-size: 18px;
    text-align: center;
    padding: 12px;
    margin-bottom: 16px;
}

/* Prompt input field */
#prompt {
    flex-grow: 1;
    border: 2px solid transparent;
    background: linear-gradient(white, white) padding-box,
                linear-gradient(135deg, #a5b4fc 0%, #e9d5ff 100%) border-box;
    border-radius: 12px;
    padding: 12px 16px;
    font-size: 16px;
    transition: all 0.3s ease;
}

#prompt:focus {
    box-shadow: 0 0 0 4px rgba(165, 180, 252, 0.25);
}

/* Run button with animated gradient */
#run_button {
    background: linear-gradient(135deg, #a78bfa 0%, #818cf8 25%, #60a5fa 50%, #34d399 75%, #fbbf24 100%);
    background-size: 200% 200%;
    animation: gradient-shift 3s ease infinite;
    color: white;
    border: none;
    padding: 12px 32px;
    border-radius: 12px;
    font-weight: 600;
    font-size: 16px;
    cursor: pointer;
    transition: all 0.3s ease;
    box-shadow: 0 4px 20px rgba(167, 139, 250, 0.4);
}

#run_button:hover {
    transform: translateY(-2px);
    box-shadow: 0 6px 30px rgba(167, 139, 250, 0.6);
}

@keyframes gradient-shift {
    0% { background-position: 0% 50%; }
    50% { background-position: 100% 50%; }
    100% { background-position: 0% 50%; }
}

/* Custom LoRA card */
.custom_lora_card {
    background: linear-gradient(135deg, #fef3c7 0%, #fde68a 100%);
    border: 1px solid #fcd34d;
    border-radius: 12px;
    padding: 16px;
    margin: 12px 0;
    box-shadow: 0 4px 12px rgba(251, 191, 36, 0.2);
}

/* Result image container */
.output-image {
    border-radius: 16px;
    overflow: hidden;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.12);
    margin-top: 20px;
}

/* Accordion styling */
.accordion {
    background: rgba(249, 250, 251, 0.9);
    border-radius: 12px;
    border: 1px solid rgba(229, 231, 235, 0.8);
    margin-top: 16px;
}

/* Slider styling */
.slider-container {
    padding: 8px 0;
}

input[type="range"] {
    background: linear-gradient(to right, #e0e7ff 0%, #c7d2fe 100%);
    border-radius: 8px;
    height: 6px;
}

/* Reuse button */
button:not(#run_button) {
    background: linear-gradient(135deg, #f0abfc 0%, #c084fc 100%);
    color: white;
    border: none;
    padding: 8px 20px;
    border-radius: 8px;
    font-weight: 500;
    cursor: pointer;
    transition: all 0.3s ease;
}

button:not(#run_button):hover {
    transform: translateY(-1px);
    box-shadow: 0 4px 16px rgba(192, 132, 252, 0.4);
}

/* Title styling */
h1 {
    background: linear-gradient(135deg, #6366f1 0%, #a855f7 25%, #ec4899 50%, #f43f5e 75%, #f59e0b 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    text-align: center;
    font-size: 3.5rem;
    font-weight: 800;
    margin-bottom: 8px;
    text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.1);
}

h1 small {
    display: block;
    background: linear-gradient(135deg, #94a3b8 0%, #64748b 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    font-size: 1rem;
    font-weight: 500;
    margin-top: 8px;
}

/* Checkbox styling */
input[type="checkbox"] {
    accent-color: #8b5cf6;
}

/* Label styling */
label {
    color: #4b5563;
    font-weight: 500;
}

/* Group containers */
.gr-group {
    background: rgba(255, 255, 255, 0.7);
    border-radius: 16px;
    padding: 20px;
    border: 1px solid rgba(255, 255, 255, 0.9);
    box-shadow: 0 4px 16px rgba(0, 0, 0, 0.05);
}
"""

# Create Gradio interface
with gr.Blocks(css=css) as demo:
    gr_flux_loras = gr.State(value=flux_loras_raw)
    
    title = gr.HTML(
        """<h1>✨ Flux-Kontext FaceLORA
        <small>Transform your portraits with AI-powered style transfer 🎨</small></h1>""",
    )
    
    selected_state = gr.State(value=None)
    custom_loaded_lora = gr.State(value=None)
    
    with gr.Row(elem_id="main_app"):
        with gr.Column(scale=4, elem_id="box_column"):
            with gr.Group(elem_id="gallery_box"):
                input_image = gr.Image(label="Upload a picture of yourself", type="pil", height=300)
                
                gallery = gr.Gallery(
                    label="Pick a LoRA",
                    allow_preview=False,
                    columns=3,
                    elem_id="gallery",
                    show_share_button=False,
                    height=400
                )
                
                custom_model = gr.Textbox(
                    label="Or enter a custom HuggingFace FLUX LoRA", 
                    placeholder="e.g., username/lora-name",
                    visible=False
                )
                custom_model_card = gr.HTML(visible=False)
                custom_model_button = gr.Button("Remove custom LoRA", visible=False)
        
        with gr.Column(scale=5):
            with gr.Row():
                prompt = gr.Textbox(
                    label="Editing Prompt",
                    show_label=False,
                    lines=1,
                    max_lines=1,
                    placeholder="optional description, e.g. 'a man with glasses and a beard'",
                    elem_id="prompt"
                )
                run_button = gr.Button("Generate ✨", elem_id="run_button")
            
            result = gr.Image(label="Generated Image", interactive=False)
            reuse_button = gr.Button("🔄 Reuse this image", visible=False)
            
            with gr.Accordion("Advanced Settings", open=False):
                lora_scale = gr.Slider(
                    label="LoRA Scale",
                    minimum=0,
                    maximum=2,
                    step=0.1,
                    value=1.5,
                    info="Controls the strength of the LoRA effect"
                )
                seed = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=0,
                )
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1,
                    maximum=10,
                    step=0.1,
                    value=2.5,
                )
            
            prompt_title = gr.Markdown(
                value="### Click on a LoRA in the gallery to select it",
                visible=True,
                elem_id="selected_lora",
            )

    # Event handlers
    custom_model.input(
        fn=load_custom_lora,
        inputs=[custom_model],
        outputs=[custom_model_card, custom_model_card, custom_model_button, custom_loaded_lora, gallery, prompt_title, selected_state],
    )
    
    custom_model_button.click(
        fn=remove_custom_lora,
        outputs=[custom_model, custom_model_button, custom_model_card, custom_loaded_lora, selected_state]
    )
    
    gallery.select(
        fn=update_selection,
        inputs=[gr_flux_loras],
        outputs=[prompt_title, prompt, selected_state],
        show_progress=False
    )
    
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer_with_lora_wrapper,
        inputs=[input_image, prompt, selected_state, custom_loaded_lora, seed, randomize_seed, guidance_scale, lora_scale, gr_flux_loras],
        outputs=[result, seed, reuse_button]
    )
    
    reuse_button.click(
        fn=lambda image: image,
        inputs=[result],
        outputs=[input_image]
    )
    
    # Initialize gallery
    demo.load(
        fn=classify_gallery, 
        inputs=[gr_flux_loras], 
        outputs=[gallery, gr_flux_loras]
    )

demo.queue(default_concurrency_limit=None)
demo.launch()