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"]
            )
            


        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])