import torch
import gradio as gr
from diffusers import DiffusionPipeline
import diffusers
import numpy as np
import random
import os
import shutil
import gc
from pathlib import Path

# =========================================================
# CACHE CLEANER FUNCTIONS
# =========================================================
def clean_huggingface_cache():
    """Clean HuggingFace cache to free up disk space"""
    cache_dir = Path.home() / ".cache" / "huggingface" / "hub"
    freed_space = 0
    
    if cache_dir.exists():
        print("🧹 Cleaning HuggingFace cache...")
        for item in cache_dir.iterdir():
            if item.is_dir():
                try:
                    size = sum(f.stat().st_size for f in item.rglob('*') if f.is_file())
                    shutil.rmtree(item)
                    freed_space += size
                    print(f"   Removed: {item.name} ({size / 1e9:.2f} GB)")
                except Exception as e:
                    print(f"   Skip {item.name}: {e}")
        print(f"✅ Freed {freed_space / 1e9:.2f} GB of disk space!")
    else:
        print("⚠️ Cache directory not found")
    
    # Force garbage collection
    gc.collect()
    torch.cuda.empty_cache() if torch.cuda.is_available() else None
    return freed_space

def clean_temp_files():
    """Clean temporary Gradio files"""
    temp_dirs = [
        Path("/tmp/gradio"),
        Path.home() / ".gradio" / "cached_examples",
        Path.home() / ".gradio" / "file_cache"
    ]
    
    freed = 0
    for temp_dir in temp_dirs:
        if temp_dir.exists():
            try:
                size = sum(f.stat().st_size for f in temp_dir.rglob('*') if f.is_file())
                shutil.rmtree(temp_dir)
                freed += size
                print(f"🧹 Cleaned: {temp_dir} ({size / 1e6:.2f} MB)")
                temp_dir.mkdir(parents=True, exist_ok=True)
            except Exception as e:
                print(f"⚠️ Could not clean {temp_dir}: {e}")
    return freed

# =========================================================
# CPU OPTIMIZATION SETTINGS (SPEED FOCUS)
# =========================================================
# Use all CPU cores for maximum speed
cpu_cores = os.cpu_count()
torch.set_num_threads(cpu_cores)  # Use ALL cores
torch.set_num_interop_threads(cpu_cores // 2)

# Enable faster CPU operations
torch.set_flush_denormal(True)  # Faster denormal handling

# Use float32 (best compatibility)
TORCH_DTYPE = torch.float32

# Disable gradient computation for inference
torch.set_grad_enabled(False)

# For even faster inference (optional - may reduce quality slightly)
os.environ["OMP_NUM_THREADS"] = str(cpu_cores)
os.environ["MKL_NUM_THREADS"] = str(cpu_cores)

# =========================================================
# MODEL CONFIGURATION
# =========================================================
MAX_SEED = np.iinfo(np.int32).max

# =========================================================
# PROMPT EXAMPLES (Keep minimal for memory)
# =========================================================
prompt_examples = [
    "The shy college girl, with glasses and a tight plaid skirt, nervously approaches her professor",
    "Her skirt rose a little higher with each gentle push, a soft blush spreading across her cheeks",
    "a girl in a school uniform having her skirt pulled up by a boy",
    "Moody mature anime scene of two lovers under neon rain, sensual atmosphere",
    "Moody mature anime scene of two lovers kissing under neon rain",
    "The girl sits on the boy's lap by the window, his hands resting on her waist",
    "A girl with long, black hair is sleeping on her desk in the classroom",
    "Her elegant silk gown swayed gracefully as she approached him",
    "A woman in a business suit having her skirt lifted by a man",
    "The older woman sits on the man's lap by the fireplace",
    "A woman with glasses, lying on the bed in just her bra, spreads her legs wide",
    "A soft focus on the girl's face, eyes closed, biting her lip",
    "A woman in a blue hanbok sits on a wooden floor, gazing out of a window",
    "The couple in a wooden outdoor bath share an intimate moment",
    "A steamy shower scene, the twins embrace under warm water",
    "The teacher pins the student against the blackboard, her skirt hiked up",
    "After hours, the girl rides the teacher on the classroom floor",
    "The secretary unzips her boss's shirt in the supply closet",
    "The secretary leans back on the boss's desk, her legs wrapped around his waist",
    "On the living room couch, one twin sits astride her sister's lap",
]

# =========================================================
# DETECT DEVICE (Prioritize speed)
# =========================================================
def get_device():
    if torch.cuda.is_available():
        print("✅ CUDA GPU detected! Using GPU for ultra-fast inference.")
        return "cuda", torch.bfloat16
    elif torch.backends.mps.is_available():
        print("✅ Apple MPS detected! Using MPS acceleration.")
        return "mps", torch.float32
    else:
        print("⚠️ No GPU detected. Running on CPU with max optimizations...")
        print(f"   Using {cpu_cores} CPU cores for parallel processing!")
        return "cpu", torch.float32

device, dtype = get_device()

# =========================================================
# MODEL CACHE STATUS
# =========================================================
def get_cache_size():
    cache_dir = Path.home() / ".cache" / "huggingface" / "hub"
    if cache_dir.exists():
        total = sum(f.stat().st_size for f in cache_dir.rglob('*') if f.is_file())
        return total / 1e9
    return 0

print(f"📦 Current HuggingFace cache size: {get_cache_size():.2f} GB")

# =========================================================
# LOAD PIPELINE (FAST LOADING)
# =========================================================
print("🚀 Loading Z-Image-Turbo pipeline (optimized for speed)...")
diffusers.utils.logging.set_verbosity_error()  # Reduce logging for speed

# Load with minimal memory footprint
pipe = DiffusionPipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    torch_dtype=dtype,
    low_cpu_mem_usage=True,
    use_safetensors=True,  # Faster loading
)

# Move to device
pipe.to(device)

# Memory optimizations
if device == "cpu":
    try:
        pipe.enable_attention_slicing()
        print("✅ Attention slicing enabled - memory efficient")
    except:
        pass
    
    # Enable CPU offload for large models
    try:
        pipe.enable_model_cpu_offload()
        print("✅ Model CPU offload enabled")
    except:
        pass

print(f"✅ Pipeline loaded on: {device}")
print(f"📦 Cache size after loading: {get_cache_size():.2f} GB")

# =========================================================
# RANDOM PROMPT FUNCTION
# =========================================================
def get_random_prompt():
    return random.choice(prompt_examples)

# =========================================================
# FAST IMAGE GENERATOR (Optimized)
# =========================================================
def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed, num_images):
    if not prompt:
        raise gr.Error("Please enter a prompt.")

    if randomize_seed:
        seed = random.randint(0, 2**32 - 1)

    num_images = min(max(1, int(num_images)), 4)
    
    # Auto-reduce steps for speed if resolution is high
    if height * width > 1024 * 1024:  # > 1MP
        actual_steps = min(num_inference_steps, 6)
        if actual_steps < num_inference_steps:
            print(f"⚡ Auto-reduced steps from {num_inference_steps} to {actual_steps} for high resolution")
    else:
        actual_steps = num_inference_steps

    generator = torch.Generator(device=device).manual_seed(int(seed))

    print(f"🎨 Generating {num_images} image(s) - Steps: {actual_steps}")

    # Generate with minimal overhead
    result = pipe(
        prompt=prompt,
        height=int(height),
        width=int(width),
        num_inference_steps=actual_steps,
        guidance_scale=0.0,
        generator=generator,
        max_sequence_length=1024,
        num_images_per_prompt=num_images,
        output_type="pil",
    )

    print(f"✅ Done! Generated {len(result.images)} image(s)")
    
    # Clean memory after generation
    gc.collect()
    
    return result.images, seed

# =========================================================
# CLEAN CACHE BUTTON FUNCTION
# =========================================================
def clean_all_cache():
    hf_freed = clean_huggingface_cache()
    temp_freed = clean_temp_files()
    total_freed = hf_freed + temp_freed
    
    status = f"🧹 Cleaned {total_freed / 1e9:.2f} GB total! (HF: {hf_freed / 1e9:.2f} GB, Temp: {temp_freed / 1e6:.2f} MB)"
    print(status)
    return status, f"Current cache: {get_cache_size():.2f} GB"

# =========================================================
# CSS (Comic Style - Minimal for faster load)
# =========================================================
css = """
@import url('https://fonts.googleapis.com/css2?family=Bangers&family=Comic+Neue:wght@400;700&display=swap');

.gradio-container {
    background-color: #FEF9C3 !important;
    background-image: radial-gradient(#1F2937 1px, transparent 1px) !important;
    background-size: 20px 20px !important;
    min-height: 100vh !important;
}

.huggingface-space-header, #space-header, .space-header, [class*="space-header"],
.svelte-1ed2p3z, .space-header-badge, .header-badge, [data-testid="space-header"],
footer, .footer, .gradio-container footer, .built-with, [class*="footer"] {
    display: none !important;
    visibility: hidden !important;
    height: 0 !important;
}

.header-text h1 {
    font-family: 'Bangers', cursive !important;
    color: #1F2937 !important;
    font-size: 3.5rem !important;
    text-align: center !important;
    margin-bottom: 0.5rem !important;
    text-shadow: 4px 4px 0px #FACC15, 6px 6px 0px #1F2937 !important;
}

.subtitle {
    text-align: center !important;
    font-family: 'Comic Neue', cursive !important;
    font-size: 1.2rem !important;
    color: #1F2937 !important;
    font-weight: 700 !important;
}

.warning-box {
    background: #FEF3C7 !important;
    border: 3px solid #F59E0B !important;
    border-radius: 8px !important;
    padding: 12px 20px !important;
    margin: 10px auto 20px auto !important;
    max-width: 800px !important;
    box-shadow: 4px 4px 0px #1F2937 !important;
}

.gr-panel, .gr-box, .block {
    background: #FFFFFF !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 6px 6px 0px #1F2937 !important;
}

.gr-button-primary, button.primary {
    background: #3B82F6 !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    color: #FFFFFF !important;
    font-family: 'Bangers', cursive !important;
    font-size: 1.3rem !important;
    padding: 14px 28px !important;
    box-shadow: 5px 5px 0px #1F2937 !important;
}

.gr-button-primary:hover, button.primary:hover {
    background: #2563EB !important;
    transform: translate(-2px, -2px) !important;
    box-shadow: 7px 7px 0px #1F2937 !important;
}

.gr-button-secondary, .random-btn, .clean-btn {
    background: #EF4444 !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    color: #FFFFFF !important;
    font-family: 'Bangers', cursive !important;
    font-size: 1.1rem !important;
    box-shadow: 4px 4px 0px #1F2937 !important;
}

.gr-button-secondary:hover, .random-btn:hover, .clean-btn:hover {
    background: #DC2626 !important;
    transform: translate(-2px, -2px) !important;
    box-shadow: 6px 6px 0px #1F2937 !important;
}

.clean-btn {
    background: #10B981 !important;
}

.clean-btn:hover {
    background: #059669 !important;
}

.gr-gallery {
    border: 4px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 8px 8px 0px #1F2937 !important;
    background: #FFFFFF !important;
}

textarea, input[type="text"], input[type="number"] {
    background: #FFFFFF !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
}

::-webkit-scrollbar { width: 12px; }
::-webkit-scrollbar-track { background: #FEF9C3; border: 2px solid #1F2937; }
::-webkit-scrollbar-thumb { background: #3B82F6; border: 2px solid #1F2937; }
"""

# =========================================================
# GRADIO UI
# =========================================================
with gr.Blocks(css=css, title="Z-Image-Turbo - Fast AI Image Generator") as demo:
    # Cache status display
    cache_status = gr.Textbox(label="📦 Cache Status", value=f"Cache size: {get_cache_size():.2f} GB", interactive=False)
    
    # HOME Badge
    gr.HTML("""
        <div style="text-align: center; margin: 20px 0 10px 0;">
            <a href="https://www.humangen.ai" target="_blank">
                <img src="https://img.shields.io/static/v1?label=🏠 HOME&message=HUMANGEN.AI&color=0000ff&labelColor=ffcc00&style=for-the-badge" alt="HOME">
            </a>
        </div>
    """)
    
    # Header
    gr.Markdown("# 🖼️ Zimage Turbo ✨", elem_classes="header-text")
    gr.Markdown('<p class="subtitle">🎨 Powered by Z-Image-Turbo - Fast Text-to-Image 🚀</p>')
    
    # Warning
    gr.HTML("""
        <div class="warning-box">
            <p>⚡ FAST MODE: 8 steps default for quick generation | 💡 Lower resolution = faster generation</p>
        </div>
    """)
    
    with gr.Row(equal_height=False):
        with gr.Column(scale=1, min_width=350):
            prompt_input = gr.Textbox(label="✏️ Prompt", placeholder="Describe your image...", lines=4)
            
            with gr.Row():
                generate_button = gr.Button("✨ GENERATE! 🎨", variant="primary", size="lg")
                clean_cache_button = gr.Button("🧹 CLEAN CACHE", elem_classes="clean-btn", size="sm")
            
            with gr.Row():
                height_input = gr.Slider(512, 2048, 768, step=64, label="📐 Height")
                width_input = gr.Slider(512, 2048, 768, step=64, label="📐 Width")
            
            num_images_input = gr.Slider(1, 4, 1, step=1, label="🖼️ Number of Images")

            with gr.Accordion("⚡ Speed Settings", open=False):
                steps_slider = gr.Slider(1, 30, 8, step=1, label="🔄 Inference Steps (lower = faster)")
                seed_input = gr.Slider(0, MAX_SEED, 42, step=1, label="🌱 Seed")
                randomize_seed_checkbox = gr.Checkbox(label="🎲 Randomize Seed", value=True)
            
            random_button = gr.Button("🎲 RANDOM PROMPT!", variant="secondary", elem_classes="random-btn")
            used_seed_output = gr.Number(label="🌱 Seed Used", interactive=False)

        with gr.Column(scale=1, min_width=350):
            output_gallery = gr.Gallery(label="🎨 Generated Images", height=550, columns=2, object_fit="contain")
            gr.Markdown('<p style="text-align: center;">💡 Click to enlarge | Right-click to save</p>')

    # Event handlers
    random_button.click(fn=get_random_prompt, outputs=[prompt_input])
    
    generate_button.click(
        fn=generate_image,
        inputs=[prompt_input, height_input, width_input, steps_slider, seed_input, randomize_seed_checkbox, num_images_input],
        outputs=[output_gallery, used_seed_output],
    )
    
    clean_cache_button.click(
        fn=clean_all_cache,
        outputs=[cache_status, cache_status],
    )

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)