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
from transformers import pipeline

translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")

# Hugging Face 토큰 설정
HF_TOKEN = os.getenv("HF_TOKEN")
if HF_TOKEN is None:
    raise ValueError("HF_TOKEN environment variable is not set")

# Setup and initialization code
cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
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)

# 샘플 이미지와 프롬프트 정의
SAMPLE_IMAGES = {
    "3d1.webp": "the most famous hero according to Zhou Qi"
}

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,
    use_auth_token=HF_TOKEN
)

# Hyper-SD LoRA 로드
pipe.load_lora_weights(
    hf_hub_download(
        "ByteDance/Hyper-SD",
        "Hyper-FLUX.1-dev-8steps-lora.safetensors",
        use_auth_token=HF_TOKEN
    )
)
pipe.fuse_lora(lora_scale=0.125)
pipe.to(device="cuda", dtype=torch.bfloat16)

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

# 예시 프롬프트 정의 부분을 수정
example_prompts = [
    "A 3D Star Wars Darth Vader helmet, highly detailed metallic finish",
    "A 3D Iron Man mask with glowing eyes and metallic red-gold finish",
    "A 3D Pokemon Pikachu figure with glossy surface",
    "A 3D geometric abstract cube transforming into a sphere, metallic finish",
    "A 3D steampunk mechanical heart with brass and copper details",
    "A 3D crystal dragon with transparent iridescent scales",
    "A 3D futuristic hovering drone with neon light accents",
    "A 3D ancient Greek warrior helmet with ornate details",
    "A 3D robotic butterfly with mechanical wings and metallic finish",
    "A 3D floating magical crystal orb with internal energy swirls"
]


@spaces.GPU
def process_and_save_image(height=1024, width=1024, steps=8, scales=3.5, prompt="", seed=None):
    global pipe
    
    if seed is None:
        seed = torch.randint(0, 1000000, (1,)).item()
    
    # 한글 감지 및 번역
    def contains_korean(text):
        return any(ord('가') <= ord(c) <= ord('힣') for c in text)
    
    # 프롬프트 전처리
    if contains_korean(prompt):
        translated = translator(prompt)[0]['translation_text']
        prompt = translated
    
    formatted_prompt = f"wbgmsst, 3D, {prompt} ,white background"
    
    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
        try:
            generated_image = pipe(
                prompt=[formatted_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
        except Exception as e:
            print(f"Error in image generation: {str(e)}")
            return None

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

def process_example(prompt):
    return process_and_save_image(
        height=1024,
        width=1024,
        steps=8,
        scales=3.5,
        prompt=prompt,
        seed=get_random_seed()
    )

# example_prompts 변수와 process_example 함수 제거

# Gradio 인터페이스
with gr.Blocks(
    theme=gr.themes.Soft(),
    css="""
        .container {
            background: linear-gradient(to bottom right, #1a1a1a, #4a4a4a);
            border-radius: 20px;
            padding: 20px;
        }
        .generate-btn {
            background: linear-gradient(45deg, #2196F3, #00BCD4);
            border: none;
            color: white;
            font-weight: bold;
            border-radius: 10px;
        }
        .output-image {
            border-radius: 15px;
            box-shadow: 0 8px 16px rgba(0,0,0,0.2);
        }
        .fixed-width {
            max-width: 1024px;
            margin: auto;
        }
        .gallery-container {
            margin-top: 40px;
            padding: 20px;
            background: #f5f5f5;
            border-radius: 15px;
        }
        .gallery-title {
            text-align: center;
            margin-bottom: 20px;
            color: #333;
            font-size: 1.5rem;
        }
    """
) as demo:
    gr.HTML(
        """
        <div style="text-align: center; max-width: 800px; margin: 0 auto; padding: 20px;">
            <h1 style="font-size: 2.5rem; color: #2196F3;">3D Style Image Generator</h1>
            <p style="font-size: 1.2rem; color: #666;">Create amazing 3D-style images with AI</p>
        </div>
        """
    )
    
    with gr.Row(elem_classes="container"):
        with gr.Column(scale=3):
            prompt = gr.Textbox(
                label="Image Description",
                placeholder="Describe the 3D 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
                    )
                
                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"]):
            output = gr.Image(
                label="Generated Image",
                elem_id="output-image",
                elem_classes=["output-image", "fixed-width"],
                value="3d.webp"
            )

    # Gallery 섹션
    with gr.Row(elem_classes="gallery-container"):
        gr.HTML("<h2 class='gallery-title'>Gallery</h2>")
        gallery_html = "<div style='display: grid; grid-template-columns: repeat(auto-fill, minmax(300px, 1fr)); gap: 20px;'>"
        
        for img_file, prompt in SAMPLE_IMAGES.items():
            if os.path.exists(img_file):
                gallery_html += f"""
                <div style='
                    border: 1px solid #ddd;
                    border-radius: 10px;
                    padding: 10px;
                    background: white;
                    box-shadow: 0 4px 8px rgba(0,0,0,0.1);
                '>
                    <img src='file/{img_file}' 
                        style='width: 100%; 
                        border-radius: 8px;
                        margin-bottom: 10px;'
                    >
                    <p style='
                        margin: 5px 0;
                        font-weight: bold;
                        color: #333;
                        padding: 10px;
                    '>Prompt: {prompt}</p>
                </div>
                """
        
        gallery_html += "</div>"
        gr.HTML(gallery_html)

# 이벤트 핸들러 부분만 수정
    def generate_image(h, w, s, sc, p, sd):
        return process_and_save_image(
            height=h,
            width=w,
            steps=s,
            scales=sc,
            prompt=p,
            seed=sd
        )

    def update_random_seed():
        return gr.Number.update(value=get_random_seed())

    # 이벤트 바인딩
    generate_btn.click(
        generate_image,
        inputs={
            "h": height,
            "w": width,
            "s": steps,
            "sc": scales,
            "p": prompt,
            "sd": seed
        },
        outputs=output
    )

    randomize_seed.click(
        update_random_seed,
        inputs=None,
        outputs=seed
    )

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