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 diffusers.pipelines.stable_diffusion import safety_checker
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)

def filter_prompt(prompt):
    # 부적절한 키워드 목록
    inappropriate_keywords = [
        # 음란/성적 키워드
        "nude", "naked", "nsfw", "porn", "sex", "explicit", "adult", "xxx",
        "erotic", "sensual", "seductive", "provocative", "intimate",
        # 폭력적 키워드
        "violence", "gore", "blood", "death", "kill", "murder", "torture",
        # 기타 부적절한 키워드
        "drug", "suicide", "abuse", "hate", "discrimination"
    ]
    
    prompt_lower = prompt.lower()
    
    # 부적절한 키워드 체크
    for keyword in inappropriate_keywords:
        if keyword in prompt_lower:
            return False, "부적절한 내용이 포함된 프롬프트입니다."
            
    return True, prompt

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)

# Add safety checker
pipe.safety_checker = safety_checker.StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")

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 {
    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;
}
.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);
}
.output-image {
    width: 100% !important;
    max-width: 100% !important;
}
.contain > div {
    width: 100% !important;
    max-width: 100% !important;
}
.fixed-width {
    width: 100% !important;
    max-width: 100% !important;
}
.gallery-container::-webkit-scrollbar {
    display: none !important;
}
.gallery-container {
    -ms-overflow-style: none !important;
    scrollbar-width: none !important;
}
#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:
        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
        
        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:
        os.makedirs(gallery_path, exist_ok=True)
        
        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)))
        
        image_files.sort(key=lambda x: x[1], reverse=True)
        
        return [f[0] for f in image_files]
    except Exception as e:
        print(f"Error loading gallery: {str(e)}")
        return []

def get_random_seed():
    return torch.randint(0, 1000000, (1,)).item()

# Create Gradio interface with a left sidebar for inputs
with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
    gr.HTML('<div class="title">FLUX Image Eternity</div>')
    gr.HTML('<div style="text-align: center; margin-bottom: 2em; color: #666;">Create stunning images from your descriptions</div>')
    
    gr.HTML("""
        <div style="color: red; margin-bottom: 1em; text-align: center; padding: 10px; background: rgba(255,0,0,0.1); border-radius: 8px;">
            ⚠️ Explicit or inappropriate content cannot be generated.
        </div>
    """)
    
    with gr.Sidebar():
        # 좌측 사이드바 메뉴: 프롬프트 입력 및 고급 설정
        prompt = gr.Textbox(
            label="Image Description",
            placeholder="Describe the image you want to create...",
            lines=3
        )
        
        with gr.Accordion("Advanced Settings", open=False):
            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
                )
            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.Markdown("## Example Prompts")
        example_prompts = [
            {
                "title": "🌅 Cinematic Landscape", 
                "prompt": "A breathtaking mountain vista at golden hour with dramatic sunbeams piercing through a sea of clouds, snow-capped peaks reflecting warm light, ultra-high detail photography, and an award-winning landscape composition reminiscent of a masterpiece captured on Hasselblad."
            },
            {
                "title": "🖼️ Fantasy Portrait", 
                "prompt": "An ethereal portrait of an elven queen with flowing silver hair, adorned with shimmering luminescent crystals and an intricate crown of twisted gold and moonstone, bathed in soft, magical lighting with highly detailed facial features that evoke ancient mysticism."
            },
            {
                "title": "🌃 Cyberpunk Scene", 
                "prompt": "A neon-drenched cyberpunk street market in torrential rain, with holographic advertisements reflecting off wet asphalt, futuristic street vendors, and a dense urban atmosphere reminiscent of Blade Runner 2049, all rendered in vivid cinematic style."
            },
            {
                "title": "🎨 Abstract Art", 
                "prompt": "A vibrant abstract composition of flowing liquid colors, dynamic swirls of iridescent purples, teals, and gold, with geometric patterns emerging from chaos—an explosion of creativity that evokes modern expressionist masterpieces."
            },
            {
                "title": "🌿 Macro Nature", 
                "prompt": "An extreme macro photograph capturing a dewdrop delicately perched on a butterfly wing, revealing intricate wing scales and rainbow refractions of light, set against a naturally blurred bokeh background, like a high-end studio shot."
            },
            {
                "title": "🤖 Futuristic Cityscape", 
                "prompt": "A vibrant futuristic cityscape at night illuminated by neon lights, holographic billboards, and flying vehicles; towering skyscrapers and a bustling urban atmosphere rendered in ultra-detailed, cinematic style."
            },
            {
                "title": "🧚 Enchanted Forest", 
                "prompt": "A mystical, enchanted forest bathed in soft, ethereal light, filled with luminescent flora, magical creatures, and ancient trees draped in moss—evoking a dreamlike fairy tale ambiance with intricate details."
            },
            {
                "title": "🌌 Cosmic Explosion", 
                "prompt": "A stunning cosmic explosion in deep space, with vibrant nebulae, swirling galaxies, and a burst of radiant colors illuminating the cosmos—captured in high-resolution digital art that is both ethereal and awe-inspiring."
            }
        ]
        # 각 예제마다 버튼을 만들어 클릭 시 prompt에 해당 문구가 자동 삽입되도록 함
        for ex in example_prompts:
            btn = gr.Button(ex["title"], variant="secondary")
            # 기본 매개변수로 ex["prompt"]를 캡처하여, 버튼 클릭 시 프롬프트에 해당 텍스트를 반환
            btn.click(fn=lambda prompt_text=ex["prompt"]: prompt_text, inputs=[], outputs=prompt)
            gr.Markdown(f"*{ex['prompt']}*")
    
    with gr.Column():
        # 메인 영역: 생성된 이미지와 갤러리 표시
        output = gr.Image(
            label="Generated Image",
            elem_id="output-image",
            elem_classes=["output-image", "fixed-width"]
        )
        
        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"]
        )
        gallery.value = load_gallery()
    
    @spaces.GPU
    def process_and_save_image(height, width, steps, scales, prompt, seed):
        # 프롬프트 필터링
        is_safe, filtered_prompt = filter_prompt(prompt)
        if not is_safe:
            gr.Warning("부적절한 내용이 포함된 프롬프트입니다.")
            return None, load_gallery()
            
        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
            try:
                generated_image = pipe(
                    prompt=[filtered_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]
                
                saved_path = save_image(generated_image)
                if saved_path is None:
                    print("Warning: Failed to save generated image")
                
                return generated_image, load_gallery()
            except Exception as e:
                print(f"Error in image generation: {str(e)}")
                return None, load_gallery()
    
    def update_seed():
        return get_random_seed()

    generate_btn.click(
        process_and_save_image,
        inputs=[height, width, steps, scales, prompt, seed],
        outputs=[output, gallery]
    )
    
    randomize_seed.click(
        update_seed,
        outputs=[seed]
    )
    
    generate_btn.click(
        update_seed,
        outputs=[seed]
    )

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