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 style selected yet", gr.update(), None
    
    lora_repo = flux_loras[selected_state.index]["repo"]
    trigger_word = flux_loras[selected_state.index]["trigger_word"]
    
    updated_text = f"### ✅ Selected Style: {flux_loras[selected_state.index]['title']}"
    new_placeholder = f"Optional: Add extra details, e.g., 'a man with glasses' or 'woman with long hair'"
    
    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), "### 📌 Please click on a style from the gallery below", 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>Custom style loaded:</strong></span>
            <div style="margin-top: 8px;">
                <h4>{repo_name}</h4>
                <small>{"Using trigger word: <code><b>"+trigger_word+"</b></code>" 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 Style: {repo_name}", None
    
    except Exception as e:
        return gr.update(visible=True), f"Error: {str(e)}", gr.update(visible=False), None, gr.update(), "### 📌 Please click on a style from the gallery below", 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
    
    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")
    # Add trigger word to prompt
    trigger_word = lora_to_use["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
css = """
#main_app {
    display: flex;
    gap: 20px;
}
#box_column {
    min-width: 400px;
}
#selected_lora {
    color: #2563eb;
    font-weight: bold;
}
#prompt {
    flex-grow: 1;
}
#run_button {
    background: linear-gradient(45deg, #2563eb, #3b82f6);
    color: white;
    border: none;
    padding: 8px 16px;
    border-radius: 6px;
    font-weight: bold;
}
.custom_lora_card {
    background: #f8fafc;
    border: 1px solid #e2e8f0;
    border-radius: 8px;
    padding: 12px;
    margin: 8px 0;
}
#gallery{
    overflow: scroll !important
}
.app-intro {
    background: linear-gradient(135deg, #f6f9fc 0%, #e9f3ff 100%);
    border-radius: 12px;
    padding: 20px;
    margin-bottom: 20px;
    border: 1px solid #e1e8f0;
}
.step-card {
    background: white;
    border: 1px solid #e2e8f0;
    border-radius: 8px;
    padding: 15px;
    margin: 10px 0;
    box-shadow: 0 2px 4px rgba(0,0,0,0.05);
}
.tips-box {
    background: #fff8dc;
    border: 1px solid #ffd700;
    border-radius: 8px;
    padding: 12px;
    margin-top: 15px;
}
"""

# Create Gradio interface
with gr.Blocks(css=css, title="aistyleportrait - AI Portrait Style Transfer Master") as demo:
    gr_flux_loras = gr.State(value=flux_loras_raw)
    
    # App title and introduction
    with gr.Column():
        gr.HTML(
            """
            <div style="text-align: center; margin-bottom: 30px;">
                <h1 style="font-size: 2.5em; margin-bottom: 10px;">
                    🎨 AI Portrait Style Transfer Master
                </h1>
                <p style="font-size: 1.2em; color: #666;">
                    Transform your photos into various artistic styles while preserving facial features
                </p>
            </div>
            """
        )
        
        # Application introduction card
        with gr.Row():
            gr.HTML(
                """
                <div class="app-intro">
                    <h2 style="margin-top: 0;">✨ Welcome to AI Portrait Style Transfer</h2>
                    <p style="line-height: 1.8;">
                        This is a powerful AI tool that can transform your portrait photos into various unique artistic styles.
                        Whether it's cartoon, oil painting, sketch, or other creative styles, it preserves your facial features
                        while giving your photos a completely new artistic expression.
                    </p>
                    <div style="margin-top: 15px;">
                        <strong>🎯 Key Features:</strong>
                        <ul style="margin-top: 8px; line-height: 1.8;">
                            <li>Preserves original facial features, ensuring person remains recognizable</li>
                            <li>Multiple preset artistic styles to choose from</li>
                            <li>Supports custom descriptions for fine-tuned effects</li>
                            <li>One-click generation, simple and fast operation</li>
                        </ul>
                    </div>
                </div>
                """
            )
    
    selected_state = gr.State(value=None)
    custom_loaded_lora = gr.State(value=None)
    
    # How to use
    with gr.Accordion("📖 How to Use (Click to expand)", open=False):
        gr.HTML(
            """
            <div style="padding: 10px;">
                <h3>🚀 Quick Start:</h3>
                <div class="step-card">
                    <h4>Step 1: Upload Your Photo</h4>
                    <p>Click the upload area on the left and select a clear portrait photo. Front-facing photos work best.</p>
                </div>
                <div class="step-card">
                    <h4>Step 2: Choose an Art Style</h4>
                    <p>Browse the style gallery and click on the artistic style you like. Each style has a preview image for reference.</p>
                </div>
                <div class="step-card">
                    <h4>Step 3: Add Description (Optional)</h4>
                    <p>If needed, you can add extra descriptions in the text box, such as "wearing glasses", "smiling", etc.</p>
                </div>
                <div class="step-card">
                    <h4>Step 4: Generate Image</h4>
                    <p>Click the "Generate" button and wait for AI to process your photo. It usually takes 10-30 seconds.</p>
                </div>
                
                <div class="tips-box">
                    <h4>💡 Pro Tips:</h4>
                    <ul style="margin-top: 8px; line-height: 1.8;">
                        <li><strong>Photo Requirements:</strong> Upload clear front-facing portrait photos with good lighting and unobstructed face</li>
                        <li><strong>Style Intensity:</strong> Adjust "LoRA Scale" in Advanced Settings to control the strength of style transfer</li>
                        <li><strong>Multiple Attempts:</strong> Enable "Randomize seed" to generate different variations</li>
                        <li><strong>Save Results:</strong> Right-click on the generated image to save it locally</li>
                    </ul>
                </div>
            </div>
            """
        )
    
    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 Your Portrait Photo", 
                    type="pil", 
                    height=300,
                    elem_classes=["upload-area"]
                )
                
                gr.Markdown("### 🎨 Choose an Art Style")
                gallery = gr.Gallery(
                    label="Style Gallery (Click to select)",
                    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 style", visible=False)
        
        with gr.Column(scale=5):
            prompt_title = gr.Markdown(
                value="### 📌 Please click on a style from the gallery on the left",
                visible=True,
                elem_id="selected_lora",
            )
            
            with gr.Row():
                prompt = gr.Textbox(
                    label="Additional Description (Optional)",
                    show_label=True,
                    lines=1,
                    max_lines=1,
                    placeholder="Optional: Add extra details, e.g., 'a man with glasses' or 'woman with long hair'",
                    elem_id="prompt"
                )
                run_button = gr.Button("🎨 Generate", elem_id="run_button", variant="primary")
            
            result = gr.Image(label="Generated Result", interactive=False)
            reuse_button = gr.Button("♻️ Use this image as new input", visible=False)
            
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                lora_scale = gr.Slider(
                    label="Style Intensity",
                    minimum=0,
                    maximum=2,
                    step=0.1,
                    value=1.5,
                    info="Controls the strength of the artistic style (recommended: 1.0-1.5)"
                )
                seed = gr.Slider(
                    label="Random Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=0,
                    info="Same seed will generate the same result"
                )
                randomize_seed = gr.Checkbox(
                    label="Randomize seed (generate different effects each time)", 
                    value=True
                )
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1,
                    maximum=10,
                    step=0.1,
                    value=2.5,
                    info="Controls how closely the image follows the prompt"
                )

    # 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()