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 = "ginipick/flux-lora-eric-cat"

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():
    predefined_images = [
        "assets/cm1.webp",
        "assets/cm2.webp",
        "assets/cm3.webp",
        "assets/cm4.webp",
        "assets/cm5.webp",
        "assets/cm6.webp",
    ]
    return predefined_images

@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 plump orange tabby cat wearing a tiny chef's hat in a sunny kitchen, flour dusted across its pink nose as it playfully bats at falling cookie dough. Its round amber eyes follow each movement with intense concentration, while its fluffy tail swishes back and forth in excitement. A small white apron tied around its rotund belly adds to the adorable scene as morning sunlight streams through lace curtains, highlighting whiskers covered in cookie crumbs. [trigger]",

    "A sleek black kitten wrapped in a miniature superhero cape, perched dramatically on top of a bookshelf. Its tiny purple cape flutters in the breeze from an open window while its bright green eyes scan the room for imaginary villains. One paw is raised heroically in the air, ready for action, while its little pink toe beans peek out beneath the flowing cape fabric. A sparkly silver mask sits slightly askew across its curious face. [trigger]",

    "A fluffy white Persian cat snuggled in a patch of garden flowers, wearing delicate fairy wings made of gossamer and glitter. Surrounded by blooming daisies and butterflies, it lazily bats at a dandelion puff floating by, while its long fur ripples like silk in the gentle spring breeze. A tiny flower crown rests between its ears, slightly crooked from its playful rolling in the grass. [trigger]",

    "A mischievous Siamese kitten in the middle of a yarn shop disaster, tangled in strings of rainbow-colored wool. Its blue eyes gleam with satisfaction as it sits in the center of its chaotic masterpiece, one paw still batting at a dangling thread. A single strand of pink yarn is draped across its ears like a festive headband, while bits of purple and blue wool decorate its whiskers. [trigger]",

    "A sleepy Scottish Fold cat in cozy pajamas, curled up in a miniature human-style bed complete with tiny pillows and a patchwork quilt. Its folded ears twitch occasionally as it dreams, while its round face wears an expression of pure contentment. The striped pajamas have little fish patterns, and one fuzzy paw clutches a diminutive teddy bear as it snoozes peacefully. [trigger]",

    "A curious tuxedo kitten exploring a magical cardboard castle, wearing a tiny paper crown and royal cape. Its black and white face peeks out from various windows and doorways it has created, whiskers twitching with excitement at each new discovery. Colorful drawings decorate the castle walls, while small paw prints in paint trail behind this creative architect as it plans its next renovation. [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"> Claude Monet 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")

    # Add sample gallery section at the bottom
    gr.Markdown("### Claude Monet Style Examples")
    predefined_gallery = gr.Gallery(
        label="Sample Images", 
        columns=3,
        rows=2,
        show_label=False, 
        value=load_predefined_images()
    )

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