import gradio as gr
import cv2
import numpy as np
import insightface
from insightface.app import FaceAnalysis
from insightface.model_zoo import get_model
import urllib.request
import os
import tempfile

# Download model if not exists
MODEL_URL = "https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx"
MODEL_PATH = "inswapper_128.onnx"

if not os.path.exists(MODEL_PATH):
    print("📥 Downloading face swap model...")
    urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)

# Initialize face analysis
print("🔄 Loading face analysis models...")
face_app = FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"])
face_app.prepare(ctx_id=0, det_size=(640, 640))
face_swapper = get_model(MODEL_PATH, providers=["CPUExecutionProvider"])
print("✅ Models ready!")

def process_video(source_img, target_video, quality_choice):
    try:
        # Quality settings map
        quality_map = {
            "320p (Fastest)": {"width": 320, "fps_reduction": 3, "bitrate": "300k"},
            "480p (Balanced)": {"width": 480, "fps_reduction": 2, "bitrate": "500k"},
            "720p (Good)": {"width": 720, "fps_reduction": 1, "bitrate": "1000k"},
            "1080p (Original)": {"width": None, "fps_reduction": 1, "bitrate": "2000k"}
        }
        
        settings = quality_map[quality_choice]
        
        # Open video
        cap = cv2.VideoCapture(target_video)
        original_fps = cap.get(cv2.CAP_PROP_FPS)
        original_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        original_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        
        # Calculate new dimensions
        if settings["width"]:
            new_width = settings["width"]
            aspect = original_height / original_width
            new_height = int(new_width * aspect)
        else:
            new_width = original_width
            new_height = original_height
        
        # Calculate new fps
        new_fps = original_fps // settings["fps_reduction"]
        
        # Setup output
        temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        out = cv2.VideoWriter(temp_output, fourcc, new_fps, (new_width, new_height))
        
        # Get source face
        source_rgb = cv2.cvtColor(source_img, cv2.COLOR_BGR2RGB)
        source_faces = face_app.get(source_rgb)
        
        if len(source_faces) == 0:
            return None, "❌ No face detected in source image"
        
        source_face = source_faces[0]
        
        # Process video
        frame_count = 0
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # Resize frame
            if settings["width"]:
                frame = cv2.resize(frame, (new_width, new_height))
            
            # Convert and swap face
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            target_faces = face_app.get(frame_rgb)
            
            if len(target_faces) > 0:
                result = face_swapper.get(frame_rgb, target_faces[0], source_face, paste_back=True)
                frame_rgb = result
            
            # Convert back and write
            frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
            out.write(frame_bgr)
            
            frame_count += 1
        
        cap.release()
        out.release()
        
        return temp_output, f"✅ Video processed at {quality_choice} - {frame_count} frames"
        
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

# Create UI
with gr.Blocks(theme=gr.themes.Soft(), title="FaceSwapAll") as demo:
    gr.Markdown("# FaceSwapAll with Quality Control")
    
    with gr.Row():
        with gr.Column():
            source = gr.Image(label="Source Face", type="numpy", height=300)
        with gr.Column():
            target = gr.Video(label="Target Video", height=300)
    
    # Quality dropdown
    quality = gr.Dropdown(
        label="Video Output Quality",
        choices=["320p (Fastest)", "480p (Balanced)", "720p (Good)", "1080p (Original)"],
        value="480p (Balanced)",
        info="Lower quality = faster processing on free CPU"
    )
    
    swap_btn = gr.Button("Swap Face in Video", variant="primary", size="lg")
    
    with gr.Row():
        result = gr.Video(label="Result Video")
        status = gr.Textbox(label="Status", lines=3)
    
    swap_btn.click(
        fn=process_video,
        inputs=[source, target, quality],
        outputs=[result, status]
    )
    
    gr.Markdown("""
    ### ⚡ Speed Tips:
    - **320p**: Fastest (8-10x speedup) - Best for testing
    - **480p**: Balanced (4-5x speedup) - Good quality/speed tradeoff
    - **720p**: Good quality (2x speedup)
    - **1080p**: Original quality - Slowest on free CPU
    """)

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