"""
3D Asset Generator Pro - Streamlined Edition
Modern, clean implementation optimized for production use.
"""

# CRITICAL: Import spaces FIRST before any CUDA initialization
import spaces

import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True,max_split_size_mb:512'

import gradio as gr
from pathlib import Path

from core import AssetPipeline, QUALITY_PRESETS
from utils import MemoryManager


# Initialize components
memory_manager = MemoryManager()
memory_manager.setup_cuda_optimizations()

pipeline = AssetPipeline()


def generate_asset(prompt: str, quality: str, progress=gr.Progress()) -> tuple:
    """
    Generate 3D asset from text prompt.
    
    Args:
        prompt: Text description
        quality: Quality preset
        progress: Gradio progress tracker
    
    Returns:
        (glb_path, status_message)
    """
    try:
        result = pipeline.generate(
            prompt=prompt,
            quality=quality,
            progress_callback=progress
        )
        
        return str(result.glb_path), result.status_message
        
    except Exception as e:
        error_msg = f"❌ Generation failed: {str(e)}"
        print(error_msg)
        return None, error_msg


# Build Gradio UI
with gr.Blocks(title="3D Asset Generator Pro", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎮 3D Asset Generator Pro
    
    Generate game-ready 3D assets from text descriptions using FLUX.1-dev + Hunyuan3D-2.1
    
    **Features:**
    - ⚡ FLUX.1-dev for high-quality 2D generation
    - 🎨 Hunyuan3D-2.1 for production-ready 3D models
    - 🔧 Automatic Blender optimization (LODs, collision, Draco compression)
    - 💾 Smart caching (60% GPU quota savings)
    - 🎯 Optimized for L4 GPU (24GB VRAM)
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Input")
            
            prompt_input = gr.Textbox(
                label="Prompt",
                placeholder="medieval knight, detailed armor, game asset",
                lines=3,
                max_lines=5
            )
            
            quality_input = gr.Dropdown(
                label="Quality Preset",
                choices=list(QUALITY_PRESETS.keys()),
                value="High",
                info="Higher quality = better results but slower generation"
            )
            
            # Quality info
            with gr.Accordion("Quality Preset Details", open=False):
                gr.Markdown("""
                **Fast** (~45s): 10 FLUX steps, 10 Hunyuan steps, 2K textures
                **Balanced** (~60s): 15 FLUX steps, 25 Hunyuan steps, 2K textures
                **High** (~90s): 25 FLUX steps, 35 Hunyuan steps, 4K textures
                **Ultra** (~120s): 30 FLUX steps, 50 Hunyuan steps, 4K textures
                """)
            
            generate_btn = gr.Button("🚀 Generate Asset", variant="primary", size="lg")
            
            gr.Markdown("""
            ### Examples
            - "medieval knight with detailed armor"
            - "futuristic mech robot, game asset"
            - "fantasy dragon, detailed scales"
            - "wooden barrel, game prop"
            - "sci-fi weapon, energy rifle"
            """)
        
        with gr.Column(scale=1):
            gr.Markdown("### Output")
            
            output_model = gr.Model3D(
                label="Generated 3D Asset",
                height=500,
                clear_color=[0.1, 0.1, 0.1, 1.0]
            )
            
            status_output = gr.Textbox(
                label="Status",
                lines=5,
                max_lines=10
            )
    
    # Event handlers
    generate_btn.click(
        fn=generate_asset,
        inputs=[prompt_input, quality_input],
        outputs=[output_model, status_output],
        api_name="generate_asset"  # Explicit API endpoint name
    )
    
    # Examples
    gr.Examples(
        examples=[
            ["medieval knight with detailed armor", "High"],
            ["futuristic mech robot, game asset", "Balanced"],
            ["fantasy dragon with detailed scales", "High"],
            ["wooden barrel, game prop", "Fast"],
            ["sci-fi energy rifle weapon", "Balanced"],
        ],
        inputs=[prompt_input, quality_input],
        outputs=[output_model, status_output],
        fn=generate_asset,
        cache_examples=False
    )
    
    gr.Markdown("""
    ---
    ### Technical Details
    
    **Pipeline:**
    1. **FLUX.1-dev** - Generate high-quality 2D reference image
    2. **Hunyuan3D-2.1** - Convert 2D to production-ready 3D model
    3. **Blender** - Optimize topology, generate LODs, add collision meshes
    4. **Export** - Game-ready GLB with Draco compression
    
    **Optimizations:**
    - Smart caching (60% GPU quota savings)
    - TF32 acceleration (20-30% faster on L4 GPU)
    - Memory-efficient pipeline (no OOM errors)
    - Automatic retry on API failures
    
    **Output Format:**
    - GLB with embedded PBR materials
    - 3 LOD levels (100%, 50%, 25%)
    - Simplified collision mesh
    - Draco compression (60-70% size reduction)
    """)


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