import gradio as gr
from typing import Optional, List, Dict, Any
import numpy as np
import time
import sys
from pathlib import Path

# CRITICAL FIX: Ensure local modules are prioritized over system packages
# This prevents ImportError when a 'utils' package exists in site-packages
current_dir = Path(__file__).resolve().parent
if str(current_dir) not in sys.path:
    sys.path.insert(0, str(current_dir))

# Clear any cached utils module to force reimport from local file
if 'utils' in sys.modules:
    del sys.modules['utils']

# Now import from the local utils.py file
from utils import WANVideoGenerator, LoRAManager, NSFWChecker
from config import MODEL_CONFIGS, AVAILABLE_LORAS, NSFW_CONFIG

# Initialize core components
generator = WANVideoGenerator()
lora_manager = LoRAManager()
nsfw_checker = NSFWChecker()

def generate_video(
    image: np.ndarray,
    prompt: str,
    selected_model: str,
    enabled_loras: List[str],
    enable_nsfw: bool,
    video_length: int,
    resolution: str,
    progress=gr.Progress()
) -> tuple[str, str, Dict[str, Any]]:
    """
    Main video generation function with WAN-scale processing
    
    Args:
        image: Input image as numpy array
        prompt: Optional text prompt for video generation
        selected_model: Selected WAN model variant
        enabled_loras: List of active LoRA adapters
        enable_nsfw: Whether to allow NSFW content generation
        video_length: Target video length in frames
        resolution: Output resolution preset
        progress: Gradio progress tracker
    
    Returns:
        Tuple of (video_path, status_message, generation_metadata)
    """
    try:
        # Step 1: Validate inputs
        progress(0.1, desc="🔍 Validating inputs...")
        if image is None:
            raise gr.Error("No image provided. Please upload an image to generate video.")
        
        # Step 2: NSFW check if enabled
        if enable_nsfw and NSFW_CONFIG["require_confirmation"]:
            progress(0.15, desc="⚠️ NSFW mode active - bypassing standard filters")
        elif not enable_nsfw:
            progress(0.15, desc="🛡️ Running safety checks...")
            if nsfw_checker.check_image(image):
                raise gr.Error("Input image flagged by safety filter. Enable NSFW mode to bypass.")
        
        # Step 3: Load selected model and LoRAs
        progress(0.2, desc=f"📦 Loading {selected_model} model...")
        generator.load_model(selected_model)
        
        progress(0.3, desc=f"🔌 Activating {len(enabled_loras)} LoRA adapters...")
        active_loras = lora_manager.load_loras(enabled_loras)
        
        # Step 4: Generate video frames
        progress(0.4, desc="🎬 Generating video frames...")
        frames = []
        for i in range(video_length):
            progress(0.4 + (i / video_length) * 0.5, 
                    desc=f"Rendering frame {i+1}/{video_length}...")
            frame = generator.generate_frame(
                image=image,
                prompt=prompt,
                frame_index=i,
                total_frames=video_length,
                active_loras=active_loras
            )
            frames.append(frame)
            time.sleep(0.1)  # Simulate processing time
        
        # Step 5: Compile video
        progress(0.95, desc="🎥 Compiling final video...")
        output_path = generator.compile_video(
            frames=frames,
            resolution=resolution,
            fps=30
        )
        
        # Step 6: Prepare metadata
        metadata = {
            "model": selected_model,
            "loras": enabled_loras,
            "nsfw_mode": enable_nsfw,
            "resolution": resolution,
            "frames": video_length,
            "prompt": prompt or "No prompt provided",
            "status": "✅ Generation complete"
        }
        
        progress(1.0, desc="✅ Done!")
        return output_path, "Video generated successfully!", metadata
    
    except Exception as e:
        raise gr.Error(f"Generation failed: {str(e)}")

def update_lora_visibility(enable_nsfw: bool) -> Dict[str, Any]:
    """Update LoRA options based on NSFW mode"""
    if enable_nsfw:
        return gr.Dropdown(
            choices=list(AVAILABLE_LORAS.keys()),
            value=[],
            multiselect=True,
            label="🎨 Active LoRA Adapters (NSFW options unlocked)"
        )
    else:
        safe_loras = {k: v for k, v in AVAILABLE_LORAS.items() if not v.get("nsfw", False)}
        return gr.Dropdown(
            choices=list(safe_loras.keys()),
            value=[],
            multiselect=True,
            label="🎨 Active LoRA Adapters (Safe mode)"
        )

def create_interface():
    """Create the main Gradio interface"""
    
    with gr.Blocks() as demo:
        gr.HTML("""
        <div style='text-align: center; padding: 20px;'>
            <h1>🥊 WAN-Scale Image-to-Video Architecture 🥊</h1>
            <p>Built with anycoder - <a href='https://huggingface.co/spaces/akhaliq/anycoder' target='_blank'>View on Hugging Face</a></p>
            <p style='font-size: 1.2em; color: #666;'>Turn static images into dynamic videos with WAN foundation models</p>
        </div>
        """)
        
        # Global state
        generation_state = gr.State({"session_id": None})
        
        with gr.Row():
            # Sidebar for controls
            with gr.Sidebar(position="left", width=320):
                gr.Markdown("### ⚙️ Generation Settings")
                
                model_selector = gr.Dropdown(
                    choices=list(MODEL_CONFIGS.keys()),
                    value="wan-2.1-14b",
                    label="🤖 WAN Model",
                    info="Select foundation model variant"
                )
                
                nsfw_toggle = gr.Checkbox(
                    value=False,
                    label="🔞 Enable NSFW Content",
                    info="Bypass safety filters (requires confirmation)"
                )
                
                lora_selector = gr.Dropdown(
                    choices=[k for k, v in AVAILABLE_LORAS.items() if not v.get("nsfw", False)],
                    value=[],
                    multiselect=True,
                    label="🎨 Active LoRA Adapters",
                    info="Select style and domain adapters"
                )
                
                with gr.Accordion("📐 Video Settings", open=False):
                    video_length = gr.Slider(
                        minimum=16,
                        maximum=128,
                        value=32,
                        step=8,
                        label="Video Length (frames)"
                    )
                    
                    resolution = gr.Radio(
                        choices=["512x512", "768x768", "1024x576", "1920x1080"],
                        value="768x768",
                        label="Resolution"
                    )
                
                with gr.Accordion("🚀 Advanced Options", open=False):
                    inference_steps = gr.Slider(
                        minimum=10,
                        maximum=100,
                        value=50,
                        label="Inference Steps"
                    )
                    
                    cfg_scale = gr.Slider(
                        minimum=1.0,
                        maximum=20.0,
                        value=7.5,
                        step=0.5,
                        label="CFG Scale"
                    )
                
                # Status indicators
                model_status = gr.Label(
                    value={"Status": "Ready", "VRAM": "24GB Available"},
                    label="System Status"
                )
            
            # Main content area
            with gr.Column():
                gr.Markdown("### 📤 Input Image")
                input_image = gr.Image(
                    label="Upload Starting Frame",
                    type="numpy",
                    height=400,
                    sources=["upload", "webcam", "clipboard"]
                )
                
                gr.Markdown("### 📝 Optional Text Prompt")
                prompt_box = gr.Textbox(
                    placeholder="Describe the motion, style, or scene...",
                    label="Prompt (optional)",
                    lines=2,
                    max_lines=4
                )
                
                with gr.Row():
                    generate_btn = gr.Button(
                        "🎬 Generate Video",
                        variant="primary",
                        scale=2
                    )
                    clear_btn = gr.ClearButton(
                        components=[input_image, prompt_box],
                        value="🗑️ Clear"
                    )
                
                # Progress tracking
                progress_bar = gr.Progress()
                status_text = gr.Textbox(
                    label="Status",
                    interactive=False,
                    show_copy_button=True
                )
                
                gr.Markdown("### 📼 Output Video")
                output_video = gr.Video(
                    label="Generated Video",
                    height=400,
                    autoplay=True,
                    show_download_button=True
                )
                
                # Generation metadata
                with gr.Accordion("📊 Generation Details", open=False):
                    metadata_json = gr.JSON(
                        label="Metadata",
                        open=False
                    )
        
        # Event handlers
        nsfw_toggle.change(
            fn=update_lora_visibility,
            inputs=nsfw_toggle,
            outputs=lora_selector,
            api_visibility="private"
        )
        
        generate_btn.click(
            fn=generate_video,
            inputs=[
                input_image,
                prompt_box,
                model_selector,
                lora_selector,
                nsfw_toggle,
                video_length,
                resolution
            ],
            outputs=[
                output_video,
                status_text,
                metadata_json
            ],
            api_visibility="public",
            concurrency_limit=2  # Limit concurrent generations
        )
        
        # Update model status on selection
        model_selector.change(
            fn=lambda x: {"Status": f"Loaded {x}", "VRAM": "24GB Used"},
            inputs=model_selector,
            outputs=model_status,
            api_visibility="private"
        )
        
        # Demo load event
        demo.load(
            fn=lambda: "System initialized and ready",
            outputs=status_text,
            api_visibility="private"
        )
    
    return demo

# Create and launch the application
if __name__ == "__main__":
    demo = create_interface()
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=False,
        show_error=True,
        max_threads=4,
        theme=gr.themes.Soft(
            primary_hue="purple",
            secondary_hue="indigo",
            neutral_hue="slate",
            font=gr.themes.GoogleFont("Inter"),
            text_size="lg",
            spacing_size="lg",
            radius_size="md"
        ).set(
            button_primary_background_fill="*primary_600",
            button_primary_background_fill_hover="*primary_700",
            block_title_text_weight="600",
            block_background_fill="*neutral_50"
        ),
        footer_links=[
            {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"},
            {"label": "Model Docs", "url": "https://huggingface.co/docs"},
            {"label": "API Reference", "url": "/docs"}
        ],
        css="""
        .gradio-container { max-width: 1400px; margin: auto; }
        .contain { display: flex; flex-direction: column; height: 100vh; }
        #component-0 { height: 100%; }
        .gr-button { font-weight: 600; }
        .gr-markdown { text-align: center; }
        """
    )