import os
import gradio as gr
import torch
import spaces
from PIL import Image
import tempfile
import subprocess
import sys
import time
import shutil
from huggingface_hub import snapshot_download

# Configuration
MODEL_NAME = "Skywork/Matrix-Game-2.0"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

print(f"🚀 Matrix-Game-2.0 Clean Setup")
print(f"📱 Device: {DEVICE}")
print(f"🔥 CUDA: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"🎮 GPU: {torch.cuda.get_device_name()}")

@spaces.GPU(duration=900)  # 15 minutes
def generate_matrix_video(input_image, num_frames, seed, use_streaming):
    """
    Matrix-Game-2.0 generation following official workflow
    """
    if input_image is None:
        return None, "❌ Please upload an input image"
    
    log = ["🚀 **MATRIX-GAME-2.0 CLEAN GENERATION**\n"]
    original_cwd = os.getcwd()
    
    try:
        # Step 1: Clone repository (official workflow)
        log.append("📥 **STEP 1: git clone Matrix-Game**")
        base_dir = os.getcwd()
        matrix_root = os.path.join(base_dir, "Matrix-Game")
        
        # Clean previous installation
        if os.path.exists(matrix_root):
            shutil.rmtree(matrix_root)
            log.append("🧹 Cleaned previous installation")
        
        # Clone fresh repository
        clone_result = subprocess.run([
            'git', 'clone', 'https://github.com/SkyworkAI/Matrix-Game.git'
        ], capture_output=True, text=True, timeout=300, cwd=base_dir)
        
        if clone_result.returncode != 0:
            log.append(f"❌ Clone failed: {clone_result.stderr}")
            return None, "\n".join(log)
        
        log.append("✅ Repository cloned successfully")
        
        # Step 2: cd Matrix-Game/Matrix-Game-2 (official workflow)
        log.append("\n📂 **STEP 2: cd Matrix-Game/Matrix-Game-2**")
        matrix_2_dir = os.path.join(matrix_root, "Matrix-Game-2")
        
        if not os.path.exists(matrix_2_dir):
            log.append(f"❌ Matrix-Game-2 not found: {matrix_2_dir}")
            return None, "\n".join(log)
        
        # Change to Matrix-Game-2 directory (as per official instructions)
        os.chdir(matrix_2_dir)
        log.append(f"✅ Changed to: {os.getcwd()}")
        
        # Verify key files exist
        key_files = ['inference.py', 'requirements.txt', 'setup.py', 'configs']
        for file in key_files:
            if os.path.exists(file):
                log.append(f"✅ {file} found")
            else:
                log.append(f"❌ {file} missing")
                return None, "\n".join(log)
        
        # Step 3: pip install -r requirements.txt (official workflow)
        log.append("\n📦 **STEP 3: pip install -r requirements.txt**")
        
        req_result = subprocess.run([
            sys.executable, "-m", "pip", "install", "-r", "requirements.txt", 
            "--no-cache-dir", "--force-reinstall"
        ], capture_output=True, text=True, timeout=600)
        
        if req_result.returncode == 0:
            log.append("✅ Requirements installed successfully")
        else:
            log.append(f"⚠️ Requirements warning (continuing): {req_result.stderr[:200]}")
        
        # Step 3.5: Install missing critical dependencies
        log.append("\n🔧 **STEP 3.5: Install missing dependencies**")
        
        critical_deps = [
            "omegaconf", "einops", "transformers", "accelerate", 
            "diffusers", "opencv-python", "imageio", "imageio-ffmpeg"
        ]
        
        for dep in critical_deps:
            try:
                dep_result = subprocess.run([
                    sys.executable, "-m", "pip", "install", dep, "--no-cache-dir"
                ], capture_output=True, text=True, timeout=120)
                
                if dep_result.returncode == 0:
                    log.append(f"✅ {dep} installed")
                else:
                    log.append(f"⚠️ {dep} warning: {dep_result.stderr[:100]}")
            except Exception as e:
                log.append(f"⚠️ {dep} error: {str(e)[:100]}")
        
        # Step 4: python setup.py develop (official workflow)
        log.append("\n🔧 **STEP 4: python setup.py develop**")
        
        setup_result = subprocess.run([
            sys.executable, "setup.py", "develop"
        ], capture_output=True, text=True, timeout=300)
        
        if setup_result.returncode == 0:
            log.append("✅ Setup.py completed successfully")
        else:
            log.append(f"⚠️ Setup.py warning (continuing): {setup_result.stderr[:200]}")
        
        # Step 5: Download model weights
        log.append("\n📥 **STEP 5: Download model weights**")
        
        try:
            model_path = snapshot_download(
                repo_id=MODEL_NAME,
                cache_dir=os.path.join(base_dir, "model_cache"),
                force_download=False
            )
            log.append(f"✅ Model downloaded: {os.path.basename(model_path)}")
        except Exception as e:
            log.append(f"❌ Model download failed: {e}")
            return None, "\n".join(log)
        
        # Step 6: Prepare input image
        log.append("\n💾 **STEP 6: Prepare input image**")
        
        temp_dir = tempfile.mkdtemp(prefix="matrix_game_")
        input_path = os.path.join(temp_dir, "input.jpg")
        
        # Create outputs directory (relative path as per official instructions)
        outputs_dir = "outputs"
        if os.path.exists(outputs_dir):
            shutil.rmtree(outputs_dir)
        os.makedirs(outputs_dir, exist_ok=True)
        
        # Resize image if too large (for memory efficiency)
        original_size = input_image.size
        if max(input_image.size) > 1024:
            ratio = 1024 / max(input_image.size)
            new_size = (int(input_image.size[0] * ratio), int(input_image.size[1] * ratio))
            input_image = input_image.resize(new_size, Image.Resampling.LANCZOS)
            log.append(f"📷 Image resized: {original_size} → {input_image.size}")
        else:
            log.append(f"📷 Image size: {input_image.size}")
        
        input_image.save(input_path, "JPEG", quality=90)
        
        # Step 7: Configure inference paths
        log.append("\n🔧 **STEP 7: Configure inference**")
        
        # Find config file (relative path)
        config_dir = "configs/inference_yaml"
        config_path = None
        
        if os.path.exists(config_dir):
            yaml_files = [f for f in os.listdir(config_dir) if f.endswith(('.yaml', '.yml'))]
            if yaml_files:
                config_path = os.path.join(config_dir, yaml_files[0])
                log.append(f"✅ Config found: {config_path}")
        
        if not config_path:
            log.append(f"❌ No config found in {config_dir}")
            return None, "\n".join(log)
        
        # Find checkpoint
        checkpoint_path = None
        for root, dirs, files in os.walk(model_path):
            for file in files:
                if file.endswith(('.bin', '.pt', '.pth', '.safetensors')):
                    checkpoint_path = os.path.join(root, file)
                    break
            if checkpoint_path:
                break
        
        if not checkpoint_path:
            log.append(f"❌ No checkpoint found in {model_path}")
            return None, "\n".join(log)
        
        log.append(f"✅ Checkpoint: {os.path.basename(checkpoint_path)}")
        
        # Step 8: Run inference (official workflow)
        log.append("\n🚀 **STEP 8: Matrix-Game inference**")
        
        script_name = "inference_streaming.py" if use_streaming else "inference.py"
        
        # Build command exactly as per official instructions
        cmd = [sys.executable, script_name]
        cmd.extend([
            "--config_path", config_path,
            "--checkpoint_path", checkpoint_path,
            "--img_path", input_path,
            "--output_folder", outputs_dir,
            "--seed", str(seed),
            "--pretrained_model_path", model_path
        ])
        
        # Add num_output_frames for regular inference
        if not use_streaming:
            cmd.extend(["--num_output_frames", str(num_frames)])
        
        log.append(f"🔧 Running: python {script_name}")
        log.append(f"📂 Working directory: {os.getcwd()}")
        log.append(f"⚙️ Frames: {num_frames} | Seed: {seed} | Streaming: {use_streaming}")
        
        # Set environment for subprocess
        env = os.environ.copy()
        env['PYTHONPATH'] = matrix_2_dir
        
        try:
            # Run the inference with proper timeout
            log.append("✅ Starting Matrix-Game generation...")
            
            process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                cwd=matrix_2_dir,
                env=env
            )
            
            # Wait for completion with timeout
            try:
                stdout, stderr = process.communicate(timeout=900)  # 15 minutes
            except subprocess.TimeoutExpired:
                process.terminate()
                process.wait()
                log.append("⏰ Timeout: Generation took too long (>15 min)")
                return None, "\n".join(log)
            
            log.append(f"🔧 Process completed with code: {process.returncode}")
            
            if process.returncode != 0:
                log.append(f"❌ Inference failed:")
                log.append(f"Error: {stderr[:500]}")
                log.append(f"Output: {stdout[:200]}")
                return None, "\n".join(log)
            
            log.append("✅ Inference completed successfully!")
            
        except Exception as e:
            log.append(f"❌ Process error: {str(e)}")
            return None, "\n".join(log)
        
        # Step 9: Find generated videos
        log.append("\n📁 **STEP 9: Find generated videos**")
        
        video_files = []
        outputs_abs = os.path.join(matrix_2_dir, outputs_dir)
        
        for root, dirs, files in os.walk(outputs_abs):
            for file in files:
                if file.lower().endswith(('.mp4', '.avi', '.mov', '.mkv', '.webm')):
                    video_path = os.path.join(root, file)
                    video_files.append(video_path)
                    log.append(f"🎥 Video found: {file}")
        
        if video_files:
            final_video = video_files[0]
            file_size = os.path.getsize(final_video) / 1e6
            
            log.append(f"\n🎉 **SUCCESS!**")
            log.append(f"📊 Video size: {file_size:.1f} MB")
            log.append(f"📷 Input: {original_size}")
            log.append(f"🎮 GPU: {torch.cuda.get_device_name() if torch.cuda.is_available() else 'CPU'}")
            log.append(f"✨ Matrix-Game-2.0 generation complete!")
            
            return final_video, "\n".join(log)
        else:
            log.append("❌ No videos generated")
            
            # Debug: list all files in outputs
            if os.path.exists(outputs_abs):
                all_files = []
                for root, dirs, files in os.walk(outputs_abs):
                    for file in files:
                        all_files.append(file)
                log.append(f"📄 Files in outputs: {all_files}")
            
            return None, "\n".join(log)
        
    except Exception as e:
        log.append(f"\n❌ **CRITICAL ERROR:** {str(e)}")
        import traceback
        log.append(f"📜 Full traceback: {traceback.format_exc()}")
        return None, "\n".join(log)
    
    finally:
        # Always return to original directory
        os.chdir(original_cwd)

# Clean Gradio interface (avoiding Gradio 4.44.0 bugs)
with gr.Blocks(
    title="Matrix-Game-2.0 Clean",
    css=".container { max-width: 1200px; margin: auto; }"
) as demo:
    
    gr.HTML("""
    <div style="text-align: center; padding: 30px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; border-radius: 15px; margin-bottom: 30px;">
        <h1 style="margin: 0; font-size: 2.8em;">🎮 Matrix-Game-2.0</h1>
        <p style="margin: 15px 0; font-size: 1.3em;">Interactive World Model - Clean Implementation</p>
        <p style="margin: 0; opacity: 0.9;">⚡ Real-time generation at 25 FPS | 🎯 Precise control | 🌍 Complex environments</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 📷 Input Configuration")
            
            input_image = gr.Image(
                label="Input Image",
                type="pil",
                height=300
            )
            
            gr.Markdown("### ⚙️ Generation Settings")
            
            with gr.Row():
                num_frames = gr.Slider(
                    minimum=50,
                    maximum=300,
                    value=150,
                    step=25,
                    label="Number of Frames"
                )
                
                seed = gr.Number(
                    value=42,
                    label="Seed",
                    precision=0
                )
            
            use_streaming = gr.Checkbox(
                label="Streaming Mode",
                value=False
            )
            
            generate_btn = gr.Button(
                "🚀 Generate Matrix-Game Video",
                variant="primary",
                size="lg"
            )
            
            gr.Markdown("""
            ### 💡 Usage Tips:
            - **Upload**: Clear images with good depth and structure
            - **Frames**: 150 frames ≈ 6 seconds at 25 FPS  
            - **Time**: Generation takes 5-15 minutes depending on complexity
            - **Streaming**: Continuous generation mode (experimental)
            - **Best results**: Landscapes, cityscapes, or structured scenes
            """)
        
        with gr.Column():
            gr.Markdown("### 🎥 Generated Video")
            
            output_video = gr.Video(
                label="Matrix-Game Video Output",
                height=400
            )
            
            gr.Markdown("### 📊 Generation Log")
            
            status_log = gr.Textbox(
                label="Detailed Status and Progress",
                lines=20,
                max_lines=25,
                show_copy_button=True
            )
    
    # Connect the generation function
    generate_btn.click(
        fn=generate_matrix_video,
        inputs=[input_image, num_frames, seed, use_streaming],
        outputs=[output_video, status_log],
        show_progress=True
    )
    
    gr.HTML("""
    <div style="text-align: center; padding: 25px; margin-top: 30px; border-top: 2px solid #eee;">
        <p style="margin-bottom: 15px;">
            📖 <a href="https://arxiv.org/pdf/2508.13009" target="_blank" style="text-decoration: none;">Research Paper</a> |
            💻 <a href="https://github.com/SkyworkAI/Matrix-Game" target="_blank" style="text-decoration: none;">GitHub Repository</a> |
            🤗 <a href="https://huggingface.co/Skywork/Matrix-Game-2.0" target="_blank" style="text-decoration: none;">Model Hub</a>
        </p>
        <p style="margin: 0;"><em>⚡ Powered by Skywork AI | Clean Implementation avoiding setup issues</em></p>
    </div>
    """)

if __name__ == "__main__":
    demo.launch(share=True)