"""
Hugging Face Space for STARFlow
Text-to-Image and Text-to-Video Generation

This app allows you to run STARFlow inference on Hugging Face GPU infrastructure.
"""

# Fix OpenMP warning - MUST be set BEFORE importing torch
import os
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'
os.environ['NUMEXPR_NUM_THREADS'] = '1'
# Fix CUDA memory fragmentation
os.environ['PYTORCH_ALLOC_CONF'] = 'expandable_segments:True'

import warnings
import gradio as gr
import torch
import subprocess
import pathlib
import traceback
from pathlib import Path

# Suppress harmless warnings
warnings.filterwarnings("ignore", category=FutureWarning, message=".*torch.distributed.reduce_op.*")

# Try to import huggingface_hub for downloading checkpoints
try:
    from huggingface_hub import hf_hub_download
    HF_HUB_AVAILABLE = True
except ImportError:
    HF_HUB_AVAILABLE = False
    print("⚠️  huggingface_hub not available. Install with: pip install huggingface_hub")

# Check if running on Hugging Face Spaces
HF_SPACE = os.environ.get("SPACE_ID") is not None

# Import spaces module for ZeroGPU support (required for GPU allocation)
try:
    import spaces
    SPACES_AVAILABLE = True
except ImportError:
    SPACES_AVAILABLE = False
    print("⚠️  spaces module not available. GPU decorator will be skipped.")

# Default checkpoint paths (if uploaded to Space Files)
DEFAULT_IMAGE_CHECKPOINT = "ckpts/starflow_3B_t2i_256x256.pth"
DEFAULT_VIDEO_CHECKPOINT = "ckpts/starflow-v_7B_t2v_caus_480p_v3.pth"

# Model Hub repositories (if using Hugging Face Model Hub)
# Set these to your Model Hub repo IDs if you upload checkpoints there
# Format: "username/repo-name"
IMAGE_CHECKPOINT_REPO = "GlobalStudio/starflow-3b-checkpoint"  # Update this after creating Model Hub repo
VIDEO_CHECKPOINT_REPO = "GlobalStudio/starflow-v-7b-checkpoint"  # Update this after creating Model Hub repo

def get_checkpoint_path(checkpoint_file, default_local_path, repo_id=None, filename=None):
    """Get checkpoint path, downloading from Hub if needed."""
    # If user uploaded a file, use it
    if checkpoint_file is not None and checkpoint_file != "":
        if hasattr(checkpoint_file, 'name'):
            return checkpoint_file.name
        checkpoint_str = str(checkpoint_file)
        # If it's a file path that exists, use it
        if os.path.exists(checkpoint_str):
            return checkpoint_str
        # If it's the default path but doesn't exist, continue to download
        if checkpoint_str == default_local_path and not os.path.exists(checkpoint_str):
            pass  # Continue to download logic below
        elif checkpoint_str != default_local_path:
            return checkpoint_str
    
    # Try local path first
    if os.path.exists(default_local_path):
        return default_local_path
    
    # Try downloading from Model Hub if configured
    if repo_id and filename and HF_HUB_AVAILABLE:
        try:
            # Use /workspace if available (HF Spaces persistent), otherwise use local cache
            if os.path.exists("/workspace"):
                cache_dir = "/workspace/checkpoints"
            elif os.path.exists("/tmp"):
                cache_dir = "/tmp/checkpoints"
            else:
                # Local development: use project directory or user cache
                cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "starflow")
            os.makedirs(cache_dir, exist_ok=True)
            
            # Check if already downloaded
            possible_path = os.path.join(cache_dir, "models--" + repo_id.replace("/", "--"), "snapshots", "*", filename)
            import glob
            existing = glob.glob(possible_path)
            if existing:
                checkpoint_path = existing[0]
                print(f"✅ Using cached checkpoint: {checkpoint_path}")
                return checkpoint_path
            
            # Download with progress tracking
            import time
            start_time = time.time()
            print(f"📥 Downloading checkpoint from {repo_id}...")
            print(f"File size: ~15.5 GB - This may take 10-30 minutes")
            print(f"Cache directory: {cache_dir}")
            print(f"Progress will be shown below...")
            
            # Use tqdm for progress if available
            try:
                from tqdm import tqdm
                from huggingface_hub.utils import tqdm as hf_tqdm
                checkpoint_path = hf_hub_download(
                    repo_id=repo_id,
                    filename=filename,
                    cache_dir=cache_dir,
                    local_files_only=False,
                    resume_download=True,  # Resume if interrupted
                )
            except Exception as e:
                # Fallback if tqdm fails
                print(f"Note: Progress bar unavailable, downloading silently...")
                checkpoint_path = hf_hub_download(
                    repo_id=repo_id,
                    filename=filename,
                    cache_dir=cache_dir,
                    local_files_only=False,
                    resume_download=True,
                )
            
            elapsed = time.time() - start_time
            print(f"✅ Download completed in {elapsed/60:.1f} minutes")
            print(f"✅ Checkpoint at: {checkpoint_path}")
            return checkpoint_path
        except Exception as e:
            error_detail = str(e)
            if "404" in error_detail or "not found" in error_detail.lower():
                return None, f"Checkpoint not found in Model Hub.\n\nPlease verify:\n1. Repo exists: https://huggingface.co/{repo_id}\n2. File exists: {filename}\n3. Repo is Public (not Private)\n\nError: {error_detail}"
            return None, f"Error downloading checkpoint: {error_detail}\n\nThis may take 10-30 minutes for a 14GB file. Please wait or check your internet connection."
    
    # No checkpoint found
    return None, f"Checkpoint not found. Please upload a checkpoint file or configure Model Hub repository."

# Verify CUDA availability (will be True on HF Spaces with GPU hardware)
if torch.cuda.is_available():
    print(f"✅ CUDA available! Device: {torch.cuda.get_device_name(0)}")
    print(f"   CUDA Version: {torch.version.cuda}")
    print(f"   PyTorch Version: {torch.__version__}")
    print(f"   GPU Count: {torch.cuda.device_count()}")
    print(f"   Current Device: {torch.cuda.current_device()}")
else:
    print("⚠️  CUDA not available. Make sure GPU hardware is selected in Space settings.")
    print(f"   PyTorch Version: {torch.__version__}")

# Apply @spaces.GPU decorator if available (required for ZeroGPU)
# IMPORTANT: Decorator must be applied at module level for ZeroGPU to detect it at startup
def generate_image(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path):
    """Generate image from text prompt."""
    return _generate_image_impl(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path)

# Apply decorator if spaces module is available (ZeroGPU detection happens at import time)
if SPACES_AVAILABLE and hasattr(spaces, 'GPU'):
    generate_image = spaces.GPU(generate_image)
    print("✅ ZeroGPU decorator applied to generate_image")

def _generate_image_impl(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path):
    """Generate image from text prompt (implementation)."""
    # Get checkpoint path (from upload, local, or Model Hub)
    status_msg = ""
    
    # Handle checkpoint file (might be string from hidden Textbox)
    if checkpoint_file == DEFAULT_IMAGE_CHECKPOINT or checkpoint_file == "" or checkpoint_file is None:
        # Use Model Hub download
        checkpoint_file = None
    
    result = get_checkpoint_path(
        checkpoint_file, 
        DEFAULT_IMAGE_CHECKPOINT,
        IMAGE_CHECKPOINT_REPO,
        "starflow_3B_t2i_256x256.pth"
    )
    
    if isinstance(result, tuple) and result[0] is None:
        return None, result[1]  # Error message
    
    checkpoint_path = result
    
    # Show status
    status_msg += f"Using checkpoint: {checkpoint_path}\n"
    
    if not os.path.exists(checkpoint_path):
        return None, f"Error: Checkpoint file not found at {checkpoint_path}.\n\nPlease verify:\n1. Model Hub repo exists: {IMAGE_CHECKPOINT_REPO}\n2. File name matches: starflow_3B_t2i_256x256.pth\n3. Repo is Public (not Private)\n\nAttempting to download from Model Hub..."
    
    if not config_path or not os.path.exists(config_path):
        return None, "Error: Config file not found. Please ensure config file exists."
    
    status_msg += "Starting image generation...\n"
    status_msg += "⏱️  Timing breakdown:\n"
    status_msg += "   - Checkpoint download: 10-30 min (first time only, ~15.5 GB)\n"
    status_msg += "   - Model loading: 2-5 min (first time only)\n"
    status_msg += "   - Image generation: 1-3 min\n"
    status_msg += "   - Subsequent runs: Only generation time (~1-3 min)\n"
    
    try:
        # Create output directory (use /tmp for logs, outputs/ for images)
        # In HF Spaces, /tmp is accessible and outputs/ may not be visible in Files tab
        output_dir = Path("outputs")
        output_dir.mkdir(exist_ok=True)
        
        # Run sampling command
        cmd = [
            "python", "sample.py",
            "--model_config_path", config_path,
            "--checkpoint_path", checkpoint_path,
            "--caption", prompt,
            "--sample_batch_size", "1",
            "--cfg", str(cfg),
            "--aspect_ratio", aspect_ratio,
            "--seed", str(seed),
            "--save_folder", "1",
            "--finetuned_vae", "none",
            "--jacobi", "1",
            "--jacobi_th", "0.001",
            "--jacobi_block_size", "16",
            "--logdir", str(output_dir)  # Set logdir to outputs directory for images
        ]
        
        status_msg += "🚀 Running generation...\n"
        status_msg += "📊 Current step: Model inference (checkpoint should already be downloaded)\n"
        
        # Note about log file location
        status_msg += f"\n📋 LOGS:\n"
        status_msg += f"   All logs will be shown in the status output below\n"
        status_msg += f"   (Logs are captured in real-time)\n\n"
        
        # Ensure GPU environment variables are passed to subprocess
        env = os.environ.copy()
        # Preserve CUDA_VISIBLE_DEVICES if set (important for ZeroGPU)
        if 'CUDA_VISIBLE_DEVICES' in env:
            print(f"✅ CUDA_VISIBLE_DEVICES={env['CUDA_VISIBLE_DEVICES']}")
        
        # Verify GPU is available before starting
        if torch.cuda.is_available():
            gpu_name = torch.cuda.get_device_name(0)
            total_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
            allocated = torch.cuda.memory_allocated(0) / 1024**3
            reserved = torch.cuda.memory_reserved(0) / 1024**3
            free_memory = total_memory - reserved
            status_msg += f"✅ GPU available: {gpu_name}\n"
            status_msg += f"   Total Memory: {total_memory:.1f} GB\n"
            status_msg += f"   Allocated: {allocated:.2f} GB, Reserved: {reserved:.2f} GB, Free: {free_memory:.2f} GB\n"
            
            # Warn if memory is low
            if free_memory < 2.0:
                status_msg += f"⚠️  Warning: Low GPU memory ({free_memory:.2f} GB free). Model may not fit.\n"
        else:
            status_msg += "⚠️  Warning: CUDA not available, will use CPU (very slow)\n"
        
        # Run with timeout (45 minutes max - allows for download + generation)
        # Capture output and write to log file
        # Note: If process is killed (e.g., GPU abort), we still capture what was output
        result = subprocess.run(
            cmd, 
            capture_output=True, 
            text=True, 
            cwd=os.getcwd(),
            env=env,  # Pass environment variables (including CUDA_VISIBLE_DEVICES)
            timeout=2700
        )
        
        # Build comprehensive log content for display (not relying on file access)
        log_content_parts = []
        log_content_parts.append("=== GENERATION LOG ===\n\n")
        log_content_parts.append(f"Command: {' '.join(cmd)}\n\n")
        log_content_parts.append(f"Environment Variables:\n")
        log_content_parts.append(f"  CUDA_VISIBLE_DEVICES={env.get('CUDA_VISIBLE_DEVICES', 'not set')}\n")
        log_content_parts.append(f"  CUDA_AVAILABLE={torch.cuda.is_available()}\n")
        if torch.cuda.is_available():
            log_content_parts.append(f"  GPU_NAME={torch.cuda.get_device_name(0)}\n")
            log_content_parts.append(f"  GPU_MEMORY_TOTAL={torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB\n")
        log_content_parts.append(f"\n")
        log_content_parts.append("=== STDOUT ===\n")
        log_content_parts.append(result.stdout if result.stdout else "(empty)\n")
        log_content_parts.append("\n\n=== STDERR ===\n")
        log_content_parts.append(result.stderr if result.stderr else "(empty)\n")
        log_content_parts.append(f"\n\n=== RETURN CODE: {result.returncode} ===\n")
        
        log_content = ''.join(log_content_parts)
        
        if result.returncode != 0:
            error_msg = f"❌ Error during generation (return code: {result.returncode})\n\n"
            error_msg += f"⚠️  IMPORTANT: Scroll down in this status box to see full error details!\n"
            error_msg += f"   The error message below contains STDERR, STDOUT, and debugging info.\n\n"
            
            # Check for GPU abort or CUDA errors
            error_output = (result.stderr + result.stdout).lower()
            if "gpu aborted" in error_output or "cuda" in error_output or "out of memory" in error_output or "killed" in error_output:
                error_msg += "⚠️  GPU ERROR DETECTED\n\n"
                error_msg += "Possible causes:\n"
                error_msg += "1. GPU timeout (ZeroGPU may have a 5-10 min limit)\n"
                error_msg += "2. CUDA out of memory (model too large for GPU)\n"
                error_msg += "3. GPU allocation failed (ZeroGPU not detected)\n"
                error_msg += "4. Process killed due to memory limit\n\n"
                error_msg += "Solutions:\n"
                error_msg += "- Model is now using bfloat16/float16 for memory efficiency\n"
                error_msg += "- Try again (GPU may have been released)\n"
                error_msg += "- Check Space logs for detailed error\n"
                error_msg += "- Ensure @spaces.GPU decorator is applied\n"
                error_msg += "- Consider using paid GPU tier for longer runs\n"
                error_msg += "- If issue persists, model may be too large for available GPU\n\n"
            
            # Show detailed error information
            error_msg += f"\n{'='*80}\n"
            error_msg += f"📋 DETAILED ERROR LOGS\n"
            error_msg += f"{'='*80}\n\n"
            
            # Show return code and command
            error_msg += f"Return Code: {result.returncode}\n"
            error_msg += f"Command: {' '.join(cmd)}\n\n"
            
            # Show STDERR (usually contains the actual error)
            if result.stderr:
                error_msg += f"=== STDERR (Error Output) ===\n"
                error_msg += f"{result.stderr}\n\n"
            else:
                error_msg += f"⚠️  No STDERR output (process may have been killed silently)\n\n"
            
            # Show STDOUT (may contain useful info)
            if result.stdout:
                error_msg += f"=== STDOUT (Standard Output) ===\n"
                # Show last 5000 chars of stdout
                stdout_preview = result.stdout[-5000:] if len(result.stdout) > 5000 else result.stdout
                error_msg += f"{stdout_preview}\n\n"
            
            # Show full log content directly in error message (no file access needed)
            error_msg += f"=== FULL GENERATION LOG ===\n"
            error_msg += f"{log_content}\n\n"
            
            # Instructions on where to find more info
            error_msg += f"{'='*80}\n"
            error_msg += f"📁 ADDITIONAL DEBUGGING:\n"
            error_msg += f"{'='*80}\n"
            error_msg += f"1. Check the Space 'Logs' tab for container logs\n"
            error_msg += f"2. Look for messages from sample.py\n"
            error_msg += f"3. Check for GPU abort or CUDA errors\n"
            error_msg += f"4. All logs are shown above in this error message\n"
            error_msg += f"{'='*80}\n"
            
            return None, error_msg
        
        status_msg += "Generation complete. Looking for output...\n"
        status_msg += f"\n💡 Note: Generated images will appear directly in the UI above.\n"
        status_msg += f"   The outputs/ folder is runtime-generated and not visible in Files tab.\n\n"
        
        # Find the generated image
        # The sample.py script saves to logdir/model_name/...
        # Model name is derived from checkpoint path stem
        checkpoint_stem = Path(checkpoint_path).stem
        model_output_dir = output_dir / checkpoint_stem
        
        status_msg += f"Searching in: {model_output_dir}\n"
        status_msg += f"Also searching recursively in: {output_dir}\n"
        
        # Search in model-specific directory first, then recursively
        search_paths = [model_output_dir, output_dir]
        output_files = []
        
        for search_path in search_paths:
            if search_path.exists():
                # Look for PNG, JPG, JPEG files
                found = list(search_path.glob("**/*.png")) + list(search_path.glob("**/*.jpg")) + list(search_path.glob("**/*.jpeg"))
                output_files.extend(found)
                status_msg += f"Found {len(found)} files in {search_path}\n"
        
        if output_files:
            # Get the most recent file
            latest_file = max(output_files, key=lambda p: p.stat().st_mtime)
            # Use absolute path for Gradio to access the file
            image_path = str(latest_file.absolute())
            status_msg += f"✅ Found image: {image_path}\n"
            status_msg += f"\n💡 Note: Generated images appear directly in the UI above.\n"
            status_msg += f"   The outputs/ folder is not visible in Files tab (runtime files).\n"
            return image_path, status_msg + "\n✅ Success! Image generated."
        else:
            # Debug: list what's actually in the directory
            debug_info = f"\n\nDebug info:\n"
            debug_info += f"Output dir exists: {output_dir.exists()}\n"
            if output_dir.exists():
                debug_info += f"Contents of {output_dir}:\n"
                for item in output_dir.iterdir():
                    debug_info += f"  - {item.name} ({'dir' if item.is_dir() else 'file'})\n"
                if model_output_dir.exists():
                    debug_info += f"\nContents of {model_output_dir}:\n"
                    for item in model_output_dir.iterdir():
                        debug_info += f"  - {item.name} ({'dir' if item.is_dir() else 'file'})\n"
            
            error_msg = status_msg + f"Error: Generated image not found.\n"
            error_msg += f"Searched in: {output_dir} and {model_output_dir}\n"
            error_msg += debug_info
            if log_content:
                error_msg += f"\n\n📋 Full log details:\n{log_content[-2000:]}"
            else:
                error_msg += f"\n\nCheck stdout:\n{result.stdout[-1000:] if result.stdout else '(no output)'}"
            return None, error_msg
            
    except subprocess.TimeoutExpired:
        error_msg = f"❌ Generation timed out after 45 minutes\n\n"
        error_msg += f"This may indicate:\n"
        error_msg += f"- Model loading is taking too long\n"
        error_msg += f"- GPU timeout (ZeroGPU may have limits)\n"
        error_msg += f"- Process hung or stuck\n\n"
        error_msg += f"Try:\n"
        error_msg += f"- Check Space Logs tab for more details\n"
        error_msg += f"- Try generating again\n"
        error_msg += f"- Check if GPU is still available\n"
        return None, error_msg
    except Exception as e:
        # Get full traceback for debugging
        error_traceback = traceback.format_exc()
        error_msg = f"❌ Exception occurred during image generation:\n\n"
        error_msg += f"Error Type: {type(e).__name__}\n"
        error_msg += f"Error Message: {str(e)}\n\n"
        error_msg += f"=== FULL TRACEBACK ===\n{error_traceback}\n\n"
        error_msg += f"💡 TIP: Scroll down in this status box to see full error details.\n"
        error_msg += f"   You can also copy the error message using the copy button.\n"
        return None, error_msg


# Apply @spaces.GPU decorator if available (required for ZeroGPU)
# IMPORTANT: Decorator must be applied at module level for ZeroGPU to detect it at startup
def generate_video(prompt, aspect_ratio, cfg, seed, target_length, checkpoint_file, config_path, input_image):
    """Generate video from text prompt."""
    return _generate_video_impl(prompt, aspect_ratio, cfg, seed, target_length, checkpoint_file, config_path, input_image)

# Apply decorator if spaces module is available (ZeroGPU detection happens at import time)
if SPACES_AVAILABLE and hasattr(spaces, 'GPU'):
    generate_video = spaces.GPU(generate_video)
    print("✅ ZeroGPU decorator applied to generate_video")

def _generate_video_impl(prompt, aspect_ratio, cfg, seed, target_length, checkpoint_file, config_path, input_image):
    """Generate video from text prompt (implementation)."""
    # Handle checkpoint file (might be string from hidden Textbox)
    if checkpoint_file == DEFAULT_VIDEO_CHECKPOINT or checkpoint_file == "" or checkpoint_file is None:
        # Use Model Hub download
        checkpoint_file = None
    
    # Get checkpoint path (from upload, local, or Model Hub)
    result = get_checkpoint_path(
        checkpoint_file,
        DEFAULT_VIDEO_CHECKPOINT,
        VIDEO_CHECKPOINT_REPO,
        "starflow-v_7B_t2v_caus_480p_v3.pth"
    )
    
    if isinstance(result, tuple) and result[0] is None:
        return None, result[1]  # Error message
    
    checkpoint_path = result
    if not os.path.exists(checkpoint_path):
        return None, f"Error: Checkpoint file not found at {checkpoint_path}."
    
    if not config_path or not os.path.exists(config_path):
        return None, "Error: Config file not found. Please ensure config file exists."
    
    # Handle input image
    input_image_path = None
    if input_image is not None:
        if hasattr(input_image, 'name'):
            input_image_path = input_image.name
        else:
            input_image_path = str(input_image)
    
    try:
        # Create output directory
        output_dir = Path("outputs")
        output_dir.mkdir(exist_ok=True)
        
        # Run sampling command
        cmd = [
            "python", "sample.py",
            "--model_config_path", config_path,
            "--checkpoint_path", checkpoint_path,
            "--caption", prompt,
            "--sample_batch_size", "1",
            "--cfg", str(cfg),
            "--aspect_ratio", aspect_ratio,
            "--seed", str(seed),
            "--out_fps", "16",
            "--save_folder", "1",
            "--jacobi", "1",
            "--jacobi_th", "0.001",
            "--finetuned_vae", "none",
            "--disable_learnable_denoiser", "0",
            "--jacobi_block_size", "32",
            "--target_length", str(target_length)
        ]
        
        if input_image_path and os.path.exists(input_image_path):
            cmd.extend(["--input_image", input_image_path])
        else:
            cmd.extend(["--input_image", "none"])
        
        # Ensure GPU environment variables are passed to subprocess
        env = os.environ.copy()
        # Preserve CUDA_VISIBLE_DEVICES if set (important for ZeroGPU)
        if 'CUDA_VISIBLE_DEVICES' in env:
            print(f"✅ CUDA_VISIBLE_DEVICES={env['CUDA_VISIBLE_DEVICES']}")
        
        # Verify GPU is available before starting
        if torch.cuda.is_available():
            print(f"✅ GPU available: {torch.cuda.get_device_name(0)}")
        
        result = subprocess.run(
            cmd, 
            capture_output=True, 
            text=True, 
            cwd=os.getcwd(),
            env=env,  # Pass environment variables (including CUDA_VISIBLE_DEVICES)
            timeout=3600  # 60 minutes for video generation
        )
        
        if result.returncode != 0:
            error_msg = f"❌ Error during video generation (return code: {result.returncode})\n\n"
            
            # Check for GPU abort or CUDA errors
            error_output = (result.stderr + result.stdout).lower()
            if "gpu aborted" in error_output or "cuda" in error_output or "out of memory" in error_output:
                error_msg += "⚠️  GPU ERROR DETECTED\n\n"
                error_msg += "Possible causes:\n"
                error_msg += "1. GPU timeout (ZeroGPU may have a 5-10 min limit)\n"
                error_msg += "2. CUDA out of memory (video generation needs more GPU memory)\n"
                error_msg += "3. GPU allocation failed (ZeroGPU not detected)\n\n"
                error_msg += "Solutions:\n"
                error_msg += "- Try again (GPU may have been released)\n"
                error_msg += "- Check Space logs for detailed error\n"
                error_msg += "- Ensure @spaces.GPU decorator is applied\n"
                error_msg += "- Consider using paid GPU tier for longer runs\n\n"
            
            error_msg += f"=== STDERR ===\n{result.stderr}\n\n"
            error_msg += f"=== STDOUT ===\n{result.stdout}\n"
            return None, error_msg
        
        # Find the generated video
        output_files = list(output_dir.glob("**/*.mp4")) + list(output_dir.glob("**/*.gif"))
        if output_files:
            latest_file = max(output_files, key=lambda p: p.stat().st_mtime)
            return str(latest_file), "Success! Video generated."
        else:
            return None, "Error: Generated video not found."
            
    except Exception as e:
        # Get full traceback for debugging
        error_traceback = traceback.format_exc()
        error_msg = f"❌ Exception occurred during video generation:\n\n"
        error_msg += f"Error Type: {type(e).__name__}\n"
        error_msg += f"Error Message: {str(e)}\n\n"
        error_msg += f"=== FULL TRACEBACK ===\n{error_traceback}\n"
        return None, error_msg


# Create Gradio interface
with gr.Blocks(title="STARFlow - Text-to-Image & Video Generation") as demo:
    # Add custom CSS using gr.HTML
    gr.HTML("""
    <style>
        .gradio-container {
            font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
        }
        .main-header {
            text-align: center;
            padding: 2rem 0;
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
            border-radius: 10px;
            margin-bottom: 2rem;
        }
        .main-header h1 {
            margin: 0;
            font-size: 2.5rem;
            font-weight: 700;
            text-shadow: 2px 2px 4px rgba(0,0,0,0.2);
        }
        .main-header p {
            margin: 0.5rem 0 0 0;
            font-size: 1.1rem;
            opacity: 0.95;
        }
        .info-box {
            background: #f0f4ff;
            border-left: 4px solid #667eea;
            padding: 1rem;
            border-radius: 5px;
            margin-bottom: 1.5rem;
        }
        .input-section {
            background: #fafafa;
            padding: 1.5rem;
            border-radius: 10px;
            border: 1px solid #e0e0e0;
        }
        .output-section {
            background: white;
            padding: 1.5rem;
            border-radius: 10px;
            border: 1px solid #e0e0e0;
            box-shadow: 0 2px 8px rgba(0,0,0,0.05);
        }
        .generate-btn {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
            color: white !important;
            font-weight: 600 !important;
            padding: 0.75rem 2rem !important;
            border-radius: 8px !important;
            border: none !important;
            font-size: 1rem !important;
            transition: transform 0.2s, box-shadow 0.2s !important;
        }
        .generate-btn:hover {
            transform: translateY(-2px);
            box-shadow: 0 4px 12px rgba(102, 126, 234, 0.4) !important;
        }
        .status-box {
            font-family: 'Monaco', 'Menlo', monospace;
            font-size: 0.9rem;
            line-height: 1.6;
        }
    </style>
    """)
    # Header
    gr.HTML("""
    <div class="main-header">
        <h1>🎨 STARFlow</h1>
        <p>Scalable Transformer Auto-Regressive Flow</p>
        <p style="font-size: 0.95rem; margin-top: 0.5rem; opacity: 0.9;">
            Generate high-quality images and videos from text prompts
        </p>
    </div>
    """)
    
    # Info box
    gr.Markdown("""
    <div class="info-box">
        <strong>ℹ️ Note:</strong> Checkpoints are automatically downloaded from Model Hub on first use. 
        First generation may take 10-20 minutes for download and model loading.
    </div>
    """)
    
    with gr.Tabs() as tabs:
        with gr.Tab("🖼️ Text-to-Image", id="image_tab"):
            with gr.Row():
                with gr.Column(scale=1, min_width=400):
                    gr.Markdown("### ⚙️ Generation Settings")
                    with gr.Group():
                        image_prompt = gr.Textbox(
                            label="📝 Prompt",
                            placeholder="a film still of a cat playing piano",
                            lines=4
                        )
                        image_config = gr.Textbox(
                            label="⚙️ Config Path",
                            value="configs/starflow_3B_t2i_256x256.yaml",
                            interactive=False,
                            visible=False  # Hidden - not needed for users
                        )
                    
                    with gr.Group():
                        gr.Markdown("#### 🎨 Image Settings")
                        image_aspect = gr.Dropdown(
                            label="Aspect Ratio",
                            choices=["1:1", "2:3", "3:2", "16:9", "9:16", "4:5", "5:4"],
                            value="1:1"
                        )
                        image_cfg = gr.Slider(
                            label="CFG Scale",
                            minimum=1.0,
                            maximum=10.0,
                            value=3.6,
                            step=0.1
                        )
                        image_seed = gr.Number(
                            value=999,
                            label="🎲 Seed",
                            precision=0
                        )
                    
                    # Hidden checkpoint field
                    image_checkpoint = gr.Textbox(
                        label="Model Checkpoint Path (auto-downloaded)",
                        value=DEFAULT_IMAGE_CHECKPOINT,
                        visible=False
                    )
                    
                    image_btn = gr.Button(
                        "✨ Generate Image",
                        variant="primary",
                        size="lg",
                        elem_classes="generate-btn"
                    )
                
                with gr.Column(scale=1, min_width=500):
                    gr.Markdown("### 🎨 Generated Image")
                    image_output = gr.Image(
                        label="",
                        type="filepath",
                        height=500,
                        show_label=False
                    )
                    image_status = gr.Textbox(
                        label="📊 Status",
                        lines=15,
                        max_lines=50,
                        interactive=False,
                        elem_classes="status-box",
                        placeholder="Status messages will appear here..."
                    )
            
            image_btn.click(
                fn=generate_image,
                inputs=[image_prompt, image_aspect, image_cfg, image_seed, image_checkpoint, image_config],
                outputs=[image_output, image_status],
                show_progress=True,
                queue=True  # Use queue to handle long-running operations
            )
        
        with gr.Tab("🎬 Text-to-Video", id="video_tab"):
            with gr.Row():
                with gr.Column(scale=1, min_width=400):
                    gr.Markdown("### ⚙️ Generation Settings")
                    with gr.Group():
                        video_prompt = gr.Textbox(
                            label="📝 Prompt",
                            placeholder="a corgi dog looks at the camera",
                            lines=4
                        )
                        video_config = gr.Textbox(
                            label="⚙️ Config Path",
                            value="configs/starflow-v_7B_t2v_caus_480p.yaml",
                            interactive=False,
                            visible=False  # Hidden - not needed for users
                        )
                    
                    with gr.Group():
                        gr.Markdown("#### 🎬 Video Settings")
                        video_aspect = gr.Dropdown(
                            label="Aspect Ratio",
                            choices=["16:9", "1:1", "4:3"],
                            value="16:9"
                        )
                        video_cfg = gr.Slider(
                            label="CFG Scale",
                            minimum=1.0,
                            maximum=10.0,
                            value=3.5,
                            step=0.1
                        )
                        video_seed = gr.Number(
                            value=99,
                            label="🎲 Seed",
                            precision=0
                        )
                        video_length = gr.Slider(
                            label="Target Length (frames)",
                            minimum=81,
                            maximum=481,
                            value=81,
                            step=80
                        )
                    
                    with gr.Group():
                        gr.Markdown("#### 🖼️ Optional Input")
                        video_input_image = gr.Image(
                            label="Input Image (optional)",
                            type="filepath"
                        )
                    
                    # Hidden checkpoint field
                    video_checkpoint = gr.Textbox(
                        label="Model Checkpoint Path (auto-downloaded)",
                        value=DEFAULT_VIDEO_CHECKPOINT,
                        visible=False
                    )
                    
                    video_btn = gr.Button(
                        "✨ Generate Video",
                        variant="primary",
                        size="lg",
                        elem_classes="generate-btn"
                    )
                
                with gr.Column(scale=1, min_width=500):
                    gr.Markdown("### 🎬 Generated Video")
                    video_output = gr.Video(
                        label="",
                        height=500,
                        show_label=False
                    )
                    video_status = gr.Textbox(
                        label="📊 Status",
                        lines=15,
                        max_lines=50,
                        interactive=False,
                        elem_classes="status-box",
                        placeholder="Status messages will appear here..."
                    )
            
            video_btn.click(
                fn=generate_video,
                inputs=[video_prompt, video_aspect, video_cfg, video_seed, video_length, 
                       video_checkpoint, video_config, video_input_image],
                outputs=[video_output, video_status],
                show_progress=True,
                queue=True  # Use queue to handle long-running operations
            )

if __name__ == "__main__":
    # Enable queue for long-running operations
    demo.queue(default_concurrency_limit=1)
    
    # Password protection - users don't need HF accounts!
    # Change these to your desired username/password
    # For multiple users, use: auth=[("user1", "pass1"), ("user2", "pass2")]
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        auth=("starflow", "im30"),  # Change password!
        share=False,  # Set to True if you want public Gradio link
        max_threads=1  # Limit threads for stability
    )