# app.py - Main application file for Hugging Face Spaces
import gradio as gr
import cv2
import numpy as np
import insightface
from insightface.app import FaceAnalysis
import tempfile
import os
from moviepy.editor import VideoFileClip
import warnings
warnings.filterwarnings('ignore')

# Global variables for models
app = None
swapper = None

def initialize_models():
    """Initialize face detection and swapping models"""
    global app, swapper
    
    try:
        # Initialize face analysis
        app = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
        app.prepare(ctx_id=-1, det_size=(320, 320))
        
        # Download and load face swapper model
        model_path = "inswapper_128.onnx"
        if not os.path.exists(model_path):
            import wget
            print("Downloading face swap model...")
            wget.download("https://github.com/facefusion/facefusion-assets/releases/download/models/inswapper_128.onnx", 
                         model_path)
        
        swapper = insightface.model_zoo.get_model(model_path, download=False, download_zip=False)
        return "✅ Models loaded successfully!"
    
    except Exception as e:
        return f"❌ Error loading models: {str(e)}"

def detect_faces_in_image(image):
    """Detect and return number of faces in image"""
    if image is None:
        return "No image provided"
    
    try:
        faces = app.get(image)
        return f"✅ Detected {len(faces)} face(s) in the image"
    except Exception as e:
        return f"❌ Error detecting faces: {str(e)}"

def swap_faces_in_video(source_image, target_video, face_index=0, progress=gr.Progress()):
    """Main face swapping function"""
    
    if source_image is None or target_video is None:
        return None, "❌ Please provide both source image and target video"
    
    try:
        progress(0.1, desc="Analyzing source image...")
        
        # Extract face from source image
        faces = app.get(source_image)
        if len(faces) == 0:
            return None, "❌ No face detected in source image. Please use a clear photo with a visible face."
        
        source_face = faces[0]
        
        progress(0.2, desc="Loading video...")
        
        # Create temporary files
        temp_output = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
        temp_output_path = temp_output.name
        temp_output.close()
        
        # Process video
        cap = cv2.VideoCapture(target_video)
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        # Setup video writer
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(temp_output_path, fourcc, fps, (width, height))
        
        progress(0.3, desc="Processing frames...")
        
        frame_count = 0
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            try:
                # Detect faces in frame
                frame_faces = app.get(frame)
                
                if len(frame_faces) > face_index:
                    # Swap face
                    result_frame = swapper.get(frame, frame_faces[face_index], source_face, paste_back=True)
                    out.write(result_frame)
                else:
                    # No face to swap, use original frame
                    out.write(frame)
                
            except:
                # If frame processing fails, use original
                out.write(frame)
            
            frame_count += 1
            
            # Update progress
            if frame_count % 10 == 0:
                progress_val = 0.3 + (frame_count / total_frames) * 0.6
                progress(progress_val, desc=f"Processing frame {frame_count}/{total_frames}")
        
        cap.release()
        out.release()
        
        progress(0.9, desc="Adding audio...")
        
        # Add audio back to video
        final_output = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
        final_output_path = final_output.name
        final_output.close()
        
        try:
            video_clip = VideoFileClip(temp_output_path)
            audio_clip = VideoFileClip(target_video).audio
            
            if audio_clip is not None:
                final_clip = video_clip.set_audio(audio_clip)
                final_clip.write_videofile(final_output_path, codec='libx264', audio_codec='aac', verbose=False, logger=None)
                final_clip.close()
            else:
                # No audio in original
                os.rename(temp_output_path, final_output_path)
                
            video_clip.close()
            
        except Exception as audio_error:
            # If audio processing fails, return video without audio
            os.rename(temp_output_path, final_output_path)
        
        # Cleanup
        if os.path.exists(temp_output_path):
            os.remove(temp_output_path)
        
        progress(1.0, desc="Complete!")
        
        return final_output_path, "✅ Face swap completed successfully!"
        
    except Exception as e:
        return None, f"❌ Error during face swap: {str(e)}"

def create_interface():
    """Create the Gradio interface"""
    
    # Custom CSS for better styling
    css = """
    .gradio-container {
        max-width: 1200px;
        margin: auto;
    }
    .title {
        text-align: center;
        font-size: 2.5em;
        font-weight: bold;
        margin-bottom: 1em;
        background: linear-gradient(45deg, #ff6b6b, #4ecdc4);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
    }
    .subtitle {
        text-align: center;
        font-size: 1.2em;
        color: #666;
        margin-bottom: 2em;
    }
    """
    
    with gr.Blocks(css=css, title="AI Face Swap Studio") as demo:
        
        # Title and description
        gr.HTML("""
        <div class="title">🎭 AI Face Swap Studio</div>
        <div class="subtitle">Upload a face image and a video to swap faces using AI</div>
        """)
        
        # Initialize models on startup
        with gr.Row():
            init_btn = gr.Button("🚀 Initialize AI Models", variant="primary", size="lg")
            init_status = gr.Textbox(label="Status", value="Click to initialize models", interactive=False)
        
        init_btn.click(initialize_models, outputs=init_status)
        
        gr.Markdown("---")
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📷 Source Face Image")
                source_img = gr.Image(
                    label="Upload the face you want to use",
                    type="numpy",
                    height=300
                )
                
                # Face detection for source image
                detect_btn = gr.Button("🔍 Detect Faces", size="sm")
                face_detection_result = gr.Textbox(
                    label="Face Detection Result",
                    interactive=False
                )
                
            with gr.Column(scale=1):
                gr.Markdown("### 🎬 Target Video")
                target_vid = gr.Video(
                    label="Upload the video where you want to swap faces",
                    height=300
                )
        
        with gr.Row():
            gr.Markdown("### ⚙️ Settings")
            
        with gr.Row():
            face_index = gr.Slider(
                minimum=0,
                maximum=5,
                value=0,
                step=1,
                label="Face Index (which face to swap if multiple faces)",
                info="0 = first face, 1 = second face, etc."
            )
        
        # Process button
        with gr.Row():
            process_btn = gr.Button("🎭 Start Face Swap", variant="primary", size="lg")
        
        # Output
        with gr.Row():
            with gr.Column(scale=1):
                output_video = gr.Video(
                    label="🎉 Result Video",
                    height=400
                )
            with gr.Column(scale=1):
                output_status = gr.Textbox(
                    label="Processing Status",
                    lines=5,
                    interactive=False
                )
        
        # Examples
        gr.Markdown("---")
        gr.Markdown("### 💡 Tips for Best Results")
        gr.Markdown("""
        - **Source Image**: Use a clear, front-facing photo with good lighting
        - **Video Quality**: Higher resolution videos give better results
        - **Face Visibility**: Make sure faces are clearly visible in both image and video
        - **Processing Time**: Longer videos will take more time to process
        - **Multiple Faces**: Use the face index slider to choose which face to swap
        """)
        
        gr.Markdown("### 🚨 Important Notes")
        gr.Markdown("""
        - This tool is for entertainment and educational purposes only
        - Please respect privacy and obtain consent before using someone's likeness
        - Processing may take several minutes depending on video length
        - Maximum video length recommended: 30 seconds for faster processing
        """)
        
        # Event handlers
        detect_btn.click(
            detect_faces_in_image,
            inputs=source_img,
            outputs=face_detection_result
        )
        
        process_btn.click(
            swap_faces_in_video,
            inputs=[source_img, target_vid, face_index],
            outputs=[output_video, output_status]
        )
    
    return demo

# Launch the app
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )