import random
import os
import uuid
from datetime import datetime
import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import DiffusionPipeline
from PIL import Image

# Create permanent storage directory
SAVE_DIR = "saved_images"  # Gradio will handle the persistence
if not os.path.exists(SAVE_DIR):
    os.makedirs(SAVE_DIR, exist_ok=True)

device = "cuda" if torch.cuda.is_available() else "cpu"
repo_id = "black-forest-labs/FLUX.1-dev"
adapter_id = "seawolf2357/hanbok"

pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
pipeline.load_lora_weights(adapter_id)
pipeline = pipeline.to(device)

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

def save_generated_image(image, prompt):
    # Generate unique filename with timestamp
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    unique_id = str(uuid.uuid4())[:8]
    filename = f"{timestamp}_{unique_id}.png"
    filepath = os.path.join(SAVE_DIR, filename)
    
    # Save the image
    image.save(filepath)
    
    # Save metadata
    metadata_file = os.path.join(SAVE_DIR, "metadata.txt")
    with open(metadata_file, "a", encoding="utf-8") as f:
        f.write(f"{filename}|{prompt}|{timestamp}\n")
    
    return filepath

def load_generated_images():
    if not os.path.exists(SAVE_DIR):
        return []
    
    # Load all images from the directory
    image_files = [os.path.join(SAVE_DIR, f) for f in os.listdir(SAVE_DIR) 
                  if f.endswith(('.png', '.jpg', '.jpeg', '.webp'))]
    # Sort by creation time (newest first)
    image_files.sort(key=lambda x: os.path.getctime(x), reverse=True)
    return image_files

def load_predefined_images():
    # Return empty list since we're not using predefined images
    return []

@spaces.GPU(duration=120)
def inference(
    prompt: str,
    seed: int,
    randomize_seed: bool,
    width: int,
    height: int,
    guidance_scale: float,
    num_inference_steps: int,
    lora_scale: float,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(seed)
    
    image = pipeline(
        prompt=prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator,
        joint_attention_kwargs={"scale": lora_scale},
    ).images[0]
    
    # Save the generated image
    filepath = save_generated_image(image, prompt)
    
    # Return the image, seed, and updated gallery
    return image, seed, load_generated_images()

examples = [
    "A young woman in a silk hanbok beneath cherry blossoms, wearing a jeogori in delicate peach silk with white diagonal lines, paired with a flowing chima in jade green. Her carefully wrapped hair is adorned with a silver binyeo catching the spring light. The hanbok's gentle movement mirrors the falling cherry petals, while the contrast between the soft peach top and vibrant green skirt creates a harmonious balance with the spring scenery. [trigger]",
    
    "A noble lady in traditional hanbok beside a palace pond, dressed in a royal blue jeogori adorned with intricate gold embroidery of phoenixes. Her deep crimson chima sweeps gracefully across the weathered stone steps, its silk surface rippling like water in the autumn breeze. Traditional jewelry including a jade daenggi and gold norigae complete her elegant ensemble, catching the late afternoon sun. [trigger]",
    
    "A bride in ceremonial wedding hanbok under moonlight, wearing a dynasty red jeogori embellished with golden symmetric patterns. Her voluminous purple chima, adorned with delicate silver-threaded flowers, creates a regal silhouette against the traditional wooden pavilion. The white collar and golden tasseled norigae sway gently as she moves, her face serene beneath the traditional jokduri headdress. [trigger]",
    
    "A young musician in performance hanbok by a traditional gayageum, dressed in a snow-white jeogori with delicate embroidered plum blossoms. Her flowing indigo chima pools around her as she sits, the fabric's sheen changing with each movement. A single orchid ornament in her classic braid reflects her artistic refinement, while the hanbok's clean lines emphasize her graceful posture. [trigger]",
    
    "A court lady in winter hanbok amid falling snow, wearing a deep forest green jeogori lined with fur at the cuffs. Her burgundy chima, padded for warmth, creates elegant swirls in the snow as she walks. Delicate golden-threaded patterns at her shoulders catch the winter light, while her carved jade binyeo holds her hair in a traditional style, completing the picture of winter elegance. [trigger]",
    
    "A dancer in festival hanbok at sunset, wearing a vibrant yellow jeogori that seems to capture the golden hour light. Her twirling sapphire blue chima creates a mesmerizing display of movement and color, its silk surface reflecting the warm evening glow. Silver bangles and a coral daenggi add flashes of brilliance as she moves, the hanbok's traditional silhouette transformed into a dynamic celebration of color and motion. [trigger]"
]

css = """
footer {
    visibility: hidden;
}
"""

with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css, analytics_enabled=False) as demo:
    gr.HTML('<div class="title"> KOREA HANBOK STUDIO </div>')
    gr.HTML('<div class="title">😄Image to Video Explore: <a href="https://huggingface.co/spaces/ginigen/theater" target="_blank">https://huggingface.co/spaces/ginigen/theater</a></div>')

    with gr.Tabs() as tabs:
        with gr.Tab("Generation"):
            with gr.Column(elem_id="col-container"):
                with gr.Row():
                    prompt = gr.Text(
                        label="Prompt",
                        show_label=False,
                        max_lines=1,
                        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):
                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=42,
                    )
                    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=768,
                        )

                    with gr.Row():
                        guidance_scale = gr.Slider(
                            label="Guidance scale",
                            minimum=0.0,
                            maximum=10.0,
                            step=0.1,
                            value=3.5,
                        )
                        num_inference_steps = gr.Slider(
                            label="Number of inference steps",
                            minimum=1,
                            maximum=50,
                            step=1,
                            value=30,
                        )
                        lora_scale = gr.Slider(
                            label="LoRA scale",
                            minimum=0.0,
                            maximum=1.0,
                            step=0.1,
                            value=1.0,
                        )

                gr.Examples(
                    examples=examples,
                    inputs=[prompt],
                    outputs=[result, seed],
                )

        with gr.Tab("Gallery"):
            gallery_header = gr.Markdown("### Generated Images Gallery")
            generated_gallery = gr.Gallery(
                label="Generated Images",
                columns=6,
                show_label=False,
                value=load_generated_images(),
                elem_id="generated_gallery",
                height="auto"
            )
            refresh_btn = gr.Button("🔄 Refresh Gallery")


    # Event handlers
    def refresh_gallery():
        return load_generated_images()

    refresh_btn.click(
        fn=refresh_gallery,
        inputs=None,
        outputs=generated_gallery,
    )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=inference,
        inputs=[
            prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            lora_scale,
        ],
        outputs=[result, seed, generated_gallery],
    )

demo.queue()
demo.launch()