import streamlit as st
import numpy as np
import cv2
import insightface
from insightface.app import FaceAnalysis
import tempfile
import os

# ==================== Load Models ====================
@st.cache_resource
def load_models():
    app = FaceAnalysis(name='buffalo_l')
    app.prepare(ctx_id=0, det_size=(640, 640))
    swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=False, download_zip=False)
    return app, swapper

app, swapper = load_models()

# ==================== Face Swap Function ====================
def swap_faces_in_video(source_image, video_path, progress):
    source_faces = app.get(source_image)
    if len(source_faces) == 0:
        st.error("No face detected in the source image.")
        return None
   
    source_face = source_faces[0]
    output_path = "swapped_output.mp4"
    
    cap = cv2.VideoCapture(video_path)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
    
    if not out.isOpened():
        st.error("Failed to initialize video writer.")
        cap.release()
        return None

    for i in range(frame_count):
        ret, frame = cap.read()
        if not ret:
            break
        target_faces = app.get(frame)
        result_frame = frame.copy()
        for target_face in target_faces:
            result_frame = swapper.get(result_frame, target_face, source_face, paste_back=True)
        out.write(result_frame)
        progress.progress((i + 1) / frame_count)

    cap.release()
    out.release()
    return output_path


# ==================== Streamlit UI ====================
st.title("🎥 Face Swapper in Video")
st.write("Upload an image and a video to swap faces.")

image_file = st.file_uploader("Upload Source Image", type=["jpg", "jpeg", "png"])
video_file = st.file_uploader("Upload Video", type=["mp4", "avi", "mov"])

if st.button("🚀 Swap Faces", type="primary"):
    if image_file is not None and video_file is not None:
        try:
            # Read image
            source_image = cv2.imdecode(np.frombuffer(image_file.read(), np.uint8), cv2.IMREAD_COLOR)
           
            # Save input video temporarily
            with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_video:
                tmp_video.write(video_file.read())
                tmp_video_path = tmp_video.name

            with st.spinner("Processing video... (This can take time depending on video length)"):
                progress_bar = st.progress(0)
                output_path = swap_faces_in_video(source_image, tmp_video_path, progress_bar)

            if output_path and os.path.exists(output_path):
                st.success("✅ Face swapping completed! Video is downloading automatically...")

                # Read as bytes
                with open(output_path, "rb") as f:
                    video_bytes = f.read()

                # Display video
                st.video(video_bytes)

                # Create download button (hidden)
                download_button = st.download_button(
                    label="📥 Download Swapped Video",
                    data=video_bytes,
                    file_name="swapped_video.mp4",
                    mime="video/mp4",
                    key="auto_download"
                )

                # Auto trigger download using JavaScript
                st.components.v1.html(
                    f"""
                    <script>
                        window.onload = function() {{
                            const downloadBtn = document.querySelector('button[key="auto_download"]');
                            if (downloadBtn) {{
                                downloadBtn.click();
                            }}
                        }}
                    </script>
                    """,
                    height=0
                )

                # Cleanup
                if os.path.exists(tmp_video_path):
                    os.remove(tmp_video_path)
                if os.path.exists(output_path):
                    os.remove(output_path)

        except Exception as e:
            st.error(f"Error: {str(e)}")
    else:
        st.error("Please upload both an image and a video.")