import gradio as gr
import spaces
import torch
import cv2
import numpy as np
from PIL import Image
import os
from models import load_model, process_video_character_replacement
from utils import validate_inputs, save_uploaded_files, create_output_directory
import tempfile

# Global model variable
model = None

def initialize_model():
    """Initialize the MoCha model"""
    global model
    if model is None:
        print("Loading MoCha model...")
        model = load_model()
        print("Model loaded successfully!")

def video_character_replacement(reference_image, source_video):
    """
    Main function to perform video character replacement using MoCha model.
    
    Args:
        reference_image: PIL Image or file path of the target character
        source_video: Uploaded video file of the original subject
    
    Returns:
        Processed video with character replacement
    """
    try:
        # Initialize model if not already loaded
        if model is None:
            initialize_model()
        
        # Validate inputs
        if reference_image is None or source_video is None:
            return None, "Please provide both a reference image and a source video."
        
        # Validate and process inputs
        ref_img, video_path = validate_inputs(reference_image, source_video)
        
        # Create output directory
        output_dir = create_output_directory()
        
        # Process the video
        output_path = process_video_character_replacement(
            model, ref_img, video_path, output_dir
        )
        
        if output_path and os.path.exists(output_path):
            return output_path, "Video character replacement completed successfully!"
        else:
            return None, "Error: Failed to process video. Please check your inputs."
            
    except Exception as e:
        return None, f"Error during processing: {str(e)}"

def clear_cache():
    """Clear temporary files and cache"""
    try:
        temp_dir = tempfile.gettempdir()
        for file in os.listdir(temp_dir):
            if file.startswith('video_char_replacement_'):
                try:
                    os.remove(os.path.join(temp_dir, file))
                except:
                    pass
        return "Cache cleared successfully!"
    except Exception as e:
        return f"Error clearing cache: {str(e)}"

# Custom CSS for better styling
css = """
.container {
    max-width: 1200px;
    margin: auto;
}

.title {
    text-align: center;
    color: #2E86AB;
    font-size: 2.5em;
    margin-bottom: 10px;
}

.description {
    text-align: center;
    color: #666;
    font-size: 1.2em;
    margin-bottom: 30px;
}

.upload-area {
    border: 2px dashed #2E86AB;
    border-radius: 10px;
    padding: 20px;
    background-color: #f8f9fa;
}

.output-area {
    border: 2px solid #28a745;
    border-radius: 10px;
    padding: 20px;
    background-color: #f8fff8;
}

.status {
    font-weight: bold;
    color: #2E86AB;
}

.error {
    color: #dc3545;
    font-weight: bold;
}

.success {
    color: #28a745;
    font-weight: bold;
}
"""

# Create the Gradio interface
with gr.Blocks(css=css, title="Video Character Replacement") as demo:
    
    # Header
    gr.HTML("""
    <div class="container">
        <h1 class="title">🎬 Video Character Replacement</h1>
        <p class="description">
            End-to-End Video Character Replacement using MoCha Model<br>
            Replace subjects in videos with reference images without structural guidance
        </p>
        <p style="text-align: center; margin-bottom: 20px;">
            <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="color: #2E86AB; text-decoration: none; font-weight: bold;">Built with anycoder</a>
        </p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.HTML('<div class="upload-area">')
            # Input components
            reference_image = gr.Image(
                label="Target Character Reference Image",
                type="pil",
                height=400,
                info="Upload a clear image of the character you want to replace with"
            )
            
            source_video = gr.Video(
                label="Source Video with Original Subject",
                height=400,
                info="Upload the video containing the subject to be replaced"
            )
            gr.HTML('</div>')
            
            with gr.Row():
                process_btn = gr.Button("🚀 Replace Character", variant="primary", size="lg")
                clear_btn = gr.Button("🗑️ Clear Cache", variant="secondary")
            
            status_text = gr.Textbox(
                label="Status",
                value="Ready to process!",
                max_lines=1
            )
        
        with gr.Column(scale=1):
            gr.HTML('<div class="output-area">')
            output_video = gr.Video(
                label="Processed Video with Character Replacement",
                height=400,
                interactive=False
            )
            
            result_message = gr.Markdown(
                value="---",
                label="Result"
            )
            gr.HTML('</div>')
    
    # Process button click event
    process_btn.click(
        fn=video_character_replacement,
        inputs=[reference_image, source_video],
        outputs=[output_video, result_message]
    )
    
    # Clear cache button
    clear_btn.click(
        fn=clear_cache,
        inputs=None,
        outputs=status_text
    )
    
    # Footer
    gr.HTML("""
    <div style="text-align: center; margin-top: 30px; padding: 20px; background-color: #f8f9fa; border-radius: 10px;">
        <h3>📋 How to Use</h3>
        <p>1. Upload a clear reference image of the target character</p>
        <p>2. Upload the source video containing the subject to replace</p>
        <p>3. Click 'Replace Character' to process the video</p>
        <p>4. Download the processed video with character replacement</p>
        <br>
        <p><strong>Note:</strong> Processing large videos may take several minutes. Please be patient!</p>
    </div>
    """)

# Launch the application
if __name__ == "__main__":
    # Initialize model on startup
    initialize_model()
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=False,
        show_api=False,
        mcp_server=True  # Enable MCP server for potential integration
    )