import random
import os
import uuid
from datetime import datetime
import gradio as gr
import gradio.themes as themes  # <-- Import the theme module
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 = "openfree/president-k-dj"

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 man playing fetch with a golden retriever in a sunny park. He wears casual weekend clothes and throws a red frisbee with joy. The dog leaps gracefully through the air, tail wagging with excitement. Warm afternoon sunlight filters through the trees, creating a peaceful scene of companionship. [president_K_DJ ]",
    "A soldier standing at attention in full military gear, holding a standard-issue rifle. His uniform is crisp and properly adorned with medals. Behind him, other soldiers march in formation during a military parade. The scene conveys discipline and duty. [president_K_DJ ]",
    "A medieval knight in gleaming armor, holding an ornate sword and shield. He stands proudly in front of a majestic castle, his cape flowing in the wind. The shield bears intricate heraldic designs, and sunlight glints off his polished armor. [president_K_DJ ]",
    "A charismatic political leader addressing a crowd from a podium. He wears a well-fitted suit and gestures confidently while speaking. The audience fills a large plaza, holding supportive banners and signs. News cameras capture the moment as he delivers his speech. [president_K_DJ ]",
    "A man enjoying a peaceful morning at home, reading a newspaper at his breakfast table. He wears comfortable home clothes and sips coffee from a favorite mug. Sunlight streams through the kitchen window, and a house plant adds a touch of nature to the cozy domestic scene. [president_K_DJ]",
    "A businessman walking confidently through a modern office building. He carries a leather briefcase and wears a tailored navy suit. Floor-to-ceiling windows reveal a cityscape behind him, and his expression shows determination and purpose. [president_K_DJ ]"
]

css = """
footer {
    visibility: hidden;
}
#col-container {
    margin-top: 10px;
}
"""

# Create a customized Soft theme for better visuals
custom_theme = themes.Soft(
    primary_hue="blue",       # Color for primary actions
    secondary_hue="gray",     # Color for secondary actions
    neutral_hue="cyan",       # Neutral color scheme
    text_size="md",           # Medium text size
    spacing_size="md",        # Medium spacing for elements
    radius_size="lg",         # Larger corner radius
    font=["Open Sans", "sans-serif"]  # Modern, clean font
)

with gr.Blocks(theme=custom_theme, css=css, analytics_enabled=False) as demo:
    gr.HTML('<div class="title" style="text-align:center; font-size:2em; margin-bottom:20px;">\
             <strong> President DJ(김대중) in KOREA </strong></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,
    )

    # tie the run_button and prompt.submit to the same function
    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()