import spaces
import argparse
import os
import time
from os import path
import shutil
from datetime import datetime
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download
import gradio as gr
import torch
from diffusers import FluxPipeline
from PIL import Image


# Setup and initialization code
cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
# Use PERSISTENT_DIR environment variable for Spaces
PERSISTENT_DIR = os.environ.get("PERSISTENT_DIR", ".")
gallery_path = path.join(PERSISTENT_DIR, "gallery")

os.environ["TRANSFORMERS_CACHE"] = cache_path
os.environ["HF_HUB_CACHE"] = cache_path
os.environ["HF_HOME"] = cache_path

torch.backends.cuda.matmul.allow_tf32 = True

# Create gallery directory if it doesn't exist
if not path.exists(gallery_path):
    os.makedirs(gallery_path, exist_ok=True)

class timer:
    def __init__(self, method_name="timed process"):
        self.method = method_name
    def __enter__(self):
        self.start = time.time()
        print(f"{self.method} starts")
    def __exit__(self, exc_type, exc_val, exc_tb):
        end = time.time()
        print(f"{self.method} took {str(round(end - self.start, 2))}s")

# Model initialization
if not path.exists(cache_path):
    os.makedirs(cache_path, exist_ok=True)

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"))
pipe.fuse_lora(lora_scale=0.125)
pipe.to(device="cuda", dtype=torch.bfloat16)

css = """
footer {display: none !important}
.gradio-container {
    max-width: 1200px;
    margin: auto;
}
.contain {
    background: rgba(255, 255, 255, 0.05);
    border-radius: 12px;
    padding: 20px;
}
.generate-btn {
    background: linear-gradient(90deg, #4B79A1 0%, #283E51 100%) !important;
    border: none !important;
    color: white !important;
}
.generate-btn:hover {
    transform: translateY(-2px);
    box-shadow: 0 5px 15px rgba(0,0,0,0.2);
}
.title {
    text-align: center;
    font-size: 2.5em;
    font-weight: bold;
    margin-bottom: 1em;
    background: linear-gradient(90deg, #4B79A1 0%, #283E51 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
}
/* Gallery specific styles */
#gallery {
    width: 100% !important;
    max-width: 100% !important;
    overflow: visible !important;
}
#gallery > div {
    width: 100% !important;
    max-width: none !important;
}
#gallery > div > div {
    width: 100% !important;
    display: grid !important;
    grid-template-columns: repeat(5, 1fr) !important;
    gap: 16px !important;
    padding: 16px !important;
}
.gallery-container {
    background: rgba(255, 255, 255, 0.05);
    border-radius: 8px;
    margin-top: 10px;
    width: 100% !important;
    box-sizing: border-box !important;
}
/* Force gallery items to maintain aspect ratio */
.gallery-item {
    width: 100% !important;
    aspect-ratio: 1 !important;
    overflow: hidden !important;
    border-radius: 4px !important;
}
.gallery-item img {
    width: 100% !important;
    height: 100% !important;
    object-fit: cover !important;
    border-radius: 4px !important;
    transition: transform 0.2s;
}
.gallery-item img:hover {
    transform: scale(1.05);
}
/* Force output image container to full width */
.output-image {
    width: 100% !important;
    max-width: 100% !important;
}
/* Force container widths */
.contain > div {
    width: 100% !important;
    max-width: 100% !important;
}
.fixed-width {
    width: 100% !important;
    max-width: 100% !important;
}
/* Remove any horizontal scrolling */
.gallery-container::-webkit-scrollbar {
    display: none !important;
}
.gallery-container {
    -ms-overflow-style: none !important;
    scrollbar-width: none !important;
}
/* Ensure consistent sizing for gallery wrapper */
#gallery > div {
    width: 100% !important;
    max-width: 100% !important;
}
#gallery > div > div {
    width: 100% !important;
    max-width: 100% !important;
}
"""
def save_image(image):
    """Save the generated image and return the path"""
    try:
        # Ensure gallery directory exists
        if not os.path.exists(gallery_path):
            try:
                os.makedirs(gallery_path, exist_ok=True)
            except Exception as e:
                print(f"Failed to create gallery directory: {str(e)}")
                return None
        
        # Generate unique filename with timestamp and random suffix
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        random_suffix = os.urandom(4).hex()
        filename = f"generated_{timestamp}_{random_suffix}.png"
        filepath = os.path.join(gallery_path, filename)
        
        try:
            if isinstance(image, Image.Image):
                image.save(filepath, "PNG", quality=100)
            else:
                image = Image.fromarray(image)
                image.save(filepath, "PNG", quality=100)
            
            if not os.path.exists(filepath):
                print(f"Warning: Failed to verify saved image at {filepath}")
                return None
                
            return filepath
        except Exception as e:
            print(f"Failed to save image: {str(e)}")
            return None
            
    except Exception as e:
        print(f"Error in save_image: {str(e)}")
        return None


def load_gallery():
    """Load all images from the gallery directory"""
    try:
        # Ensure gallery directory exists
        os.makedirs(gallery_path, exist_ok=True)
        
        # Get all image files and sort by modification time
        image_files = []
        for f in os.listdir(gallery_path):
            if f.lower().endswith(('.png', '.jpg', '.jpeg')):
                full_path = os.path.join(gallery_path, f)
                image_files.append((full_path, os.path.getmtime(full_path)))
        
        # Sort by modification time (newest first)
        image_files.sort(key=lambda x: x[1], reverse=True)
        
        # Return only the file paths
        return [f[0] for f in image_files]
    except Exception as e:
        print(f"Error loading gallery: {str(e)}")
        return []

# Create Gradio interface
with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
    gr.HTML('<div class="title">AI Image Generator</div>')
    gr.HTML('<div style="text-align: center; margin-bottom: 2em; color: #666;">Create stunning images from your descriptions</div>')

    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(
                label="Image Description",
                placeholder="Describe the image you want to create...",
                lines=3
            )
            
            with gr.Accordion("Advanced Settings", open=False):
                with gr.Row():
                    height = gr.Slider(
                        label="Height",
                        minimum=256,
                        maximum=1152,
                        step=64,
                        value=1024
                    )
                    width = gr.Slider(
                        label="Width",
                        minimum=256,
                        maximum=1152,
                        step=64,
                        value=1024
                    )
                
                with gr.Row():
                    steps = gr.Slider(
                        label="Inference Steps",
                        minimum=6,
                        maximum=25,
                        step=1,
                        value=8
                    )
                    scales = gr.Slider(
                        label="Guidance Scale",
                        minimum=0.0,
                        maximum=5.0,
                        step=0.1,
                        value=3.5
                    )
                
                def get_random_seed():
                    return torch.randint(0, 1000000, (1,)).item()
                
                seed = gr.Number(
                    label="Seed (random by default, set for reproducibility)",
                    value=get_random_seed(),
                    precision=0
                )
                
                randomize_seed = gr.Button("🎲 Randomize Seed", elem_classes=["generate-btn"])
            
            generate_btn = gr.Button(
                "✨ Generate Image",
                elem_classes=["generate-btn"]
            )
            
            gr.HTML("""
                <div style="margin-top: 1em; padding: 1em; border-radius: 8px; background: rgba(255, 255, 255, 0.05);">
                    <h4 style="margin: 0 0 0.5em 0;">Example Prompts:</h4>
                    <div style="background: rgba(75, 121, 161, 0.1); padding: 1em; border-radius: 8px; margin-bottom: 1em;">
                        <p style="font-weight: bold; margin: 0 0 0.5em 0;">🌅 Cinematic Landscape</p>
                        <p style="margin: 0; font-style: italic;">"A breathtaking mountain vista at golden hour, dramatic sunbeams piercing through clouds, snow-capped peaks reflecting warm light, ultra-high detail photography, artistically composed, award-winning landscape photo, shot on Hasselblad"</p>
                    </div>
                    <div style="background: rgba(75, 121, 161, 0.1); padding: 1em; border-radius: 8px; margin-bottom: 1em;">
                        <p style="font-weight: bold; margin: 0 0 0.5em 0;">🖼️ Fantasy Portrait</p>
                        <p style="margin: 0; font-style: italic;">"Ethereal portrait of an elven queen with flowing silver hair, adorned with luminescent crystals, intricate crown of twisted gold and moonstone, soft ethereal lighting, detailed facial features, fantasy art style, highly detailed, painted by Artgerm and Charlie Bowater"</p>
                    </div>
                    <div style="background: rgba(75, 121, 161, 0.1); padding: 1em; border-radius: 8px; margin-bottom: 1em;">
                        <p style="font-weight: bold; margin: 0 0 0.5em 0;">🌃 Cyberpunk Scene</p>
                        <p style="margin: 0; font-style: italic;">"Neon-lit cyberpunk street market in rain, holographic advertisements reflecting in puddles, street vendors with glowing cyber-augmentations, dense urban environment, atmospheric fog, cinematic lighting, inspired by Blade Runner 2049"</p>
                    </div>
                    <div style="background: rgba(75, 121, 161, 0.1); padding: 1em; border-radius: 8px; margin-bottom: 1em;">
                        <p style="font-weight: bold; margin: 0 0 0.5em 0;">🎨 Abstract Art</p>
                        <p style="margin: 0; font-style: italic;">"Vibrant abstract composition of flowing liquid colors, dynamic swirls of iridescent purples and teals, golden geometric patterns emerging from chaos, luxury art style, ultra-detailed, painted in oil on canvas, inspired by James Jean and Gustav Klimt"</p>
                    </div>
                    <div style="background: rgba(75, 121, 161, 0.1); padding: 1em; border-radius: 8px; margin-bottom: 1em;">
                        <p style="font-weight: bold; margin: 0 0 0.5em 0;">🌿 Macro Nature</p>
                        <p style="margin: 0; font-style: italic;">"Extreme macro photography of a dewdrop on a butterfly wing, rainbow light refraction, crystalline clarity, intricate wing scales visible, natural bokeh background, professional studio lighting, shot with Canon MP-E 65mm lens"</p>
                    </div>
                </div>
            """)

        with gr.Column(scale=4, elem_classes=["fixed-width"]):
            # Current generated image
            output = gr.Image(
                label="Generated Image",
                elem_id="output-image",
                elem_classes=["output-image", "fixed-width"]
            )
            
            # Gallery of generated images
            gallery = gr.Gallery(
                label="Generated Images Gallery",
                show_label=True,
                elem_id="gallery",
                columns=[4],
                rows=[2],
                height="auto",
                object_fit="cover",
                elem_classes=["gallery-container", "fixed-width"]
            )
            
            # Load existing gallery images on startup
            gallery.value = load_gallery()
    
    @spaces.GPU
    def process_and_save_image(height, width, steps, scales, prompt, seed):
        global pipe
        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
            try:
                generated_image = pipe(
                    prompt=[prompt],
                    generator=torch.Generator().manual_seed(int(seed)),
                    num_inference_steps=int(steps),
                    guidance_scale=float(scales),
                    height=int(height),
                    width=int(width),
                    max_sequence_length=256
                ).images[0]
                
                # Save the generated image
                saved_path = save_image(generated_image)
                if saved_path is None:
                    print("Warning: Failed to save generated image")
                
                # Return both the generated image and updated gallery
                return generated_image, load_gallery()
            except Exception as e:
                print(f"Error in image generation: {str(e)}")
                return None, load_gallery()
    
    # Connect the generation button to both the image output and gallery update
    def update_seed():
        return get_random_seed()

    generate_btn.click(
        process_and_save_image,
        inputs=[height, width, steps, scales, prompt, seed],
        outputs=[output, gallery]
    )
    
    # Add randomize seed button functionality
    randomize_seed.click(
        update_seed,
        outputs=[seed]
    )
    
    # Also randomize seed after each generation
    generate_btn.click(
        update_seed,
        outputs=[seed]
    )

if __name__ == "__main__":
    demo.launch(allowed_paths=[PERSISTENT_DIR])