import cv2
import numpy as np
import insightface
from insightface.app import FaceAnalysis
from insightface.app.common import Face
import threading
from gfpgan.utils import GFPGANer
from typing import Any

Frame = np.ndarray


class FaceProcessor:
    """
    调用cv的模块分解mp4，组装mp4，视频质量差 
    """
    def __init__(self, face_enhancer_model: str = 'GFPGANv1.4.pth', face_swapper_model: str = 'inswapper_128.onnx', device: str = 'cuda'):
        self.FACE_ENHANCER = None
        self.FACE_SWAPPER = None
        self.THREAD_SEMAPHORE = threading.Semaphore()
        self.THREAD_LOCK = threading.Lock()

        # Set device
        self.device = device  # 'cuda' for GPU, 'cpu' for CPU

        # Load models
        self.face_enhancer_model = face_enhancer_model
        self.face_swapper_model = face_swapper_model

        # Initialize face analysis app
        self.app = FaceAnalysis(providers=['CUDAExecutionProvider' if device == 'cuda' else 'CPUExecutionProvider'])
        self.app.prepare(ctx_id=0, det_size=(640, 640))

    def get_face_enhancer(self) -> GFPGANer:
        with self.THREAD_LOCK:
            if self.FACE_ENHANCER is None:
                self.FACE_ENHANCER = GFPGANer(
                    model_path=self.face_enhancer_model, upscale=1, device=self.device)
        return self.FACE_ENHANCER

    def get_face_swapper(self) -> Any:
        with self.THREAD_LOCK:
            if self.FACE_SWAPPER is None:
                self.FACE_SWAPPER = insightface.model_zoo.get_model(
                    self.face_swapper_model, providers=['CUDAExecutionProvider' if self.device == 'cuda' else 'CPUExecutionProvider'])
        return self.FACE_SWAPPER

    def enhance_face(self, target_face: Face, temp_frame: Frame) -> Frame:
        start_x, start_y, end_x, end_y = map(int, target_face['bbox'])
        padding_x = int((end_x - start_x) * 0.5)
        padding_y = int((end_y - start_y) * 0.5)
        start_x = max(0, start_x - padding_x)
        start_y = max(0, start_y - padding_y)
        end_x = max(0, end_x + padding_x)
        end_y = max(0, end_y + padding_y)
        temp_face = temp_frame[start_y:end_y, start_x:end_x]
        if temp_face.size:
            with self.THREAD_SEMAPHORE:
                _, _, temp_face = self.get_face_enhancer().enhance(
                    temp_face,
                    paste_back=True
                )
            temp_frame[start_y:end_y, start_x:end_x] = temp_face
        return temp_frame

    def swap_face(self, source_face: Face, target_face: Face, temp_frame: Frame) -> Frame:
        return self.get_face_swapper().get(temp_frame, target_face, source_face, paste_back=True)

    def process_video(self, src_mp4: str, image_path: str, dst_mp4: str) -> None:
        # Read the target image (source face)
        target_image = cv2.imread(image_path)
        target_faces = self.app.get(target_image)
        if len(target_faces) == 0:
            print("No face detected in the target image.")
            return
        target_face = target_faces[0]

        # Open the source video
        cap = cv2.VideoCapture(src_mp4)
        if not cap.isOpened():
            print("Error opening video stream or file")
            return

        # Get video properties
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        # Create VideoWriter object to save the output video
        out = cv2.VideoWriter(dst_mp4, fourcc, fps, (frame_width, frame_height))

        # Get the total frame count
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

        # Process each frame of the video
        frame_count = 0
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            # Detect faces in the current frame
            faces = self.app.get(frame)
            if len(faces) > 0:
                # Swap face from the target image to the current frame
                swapped_face_frame = self.swap_face(target_face, faces[0], frame)
                self.enhance_face(faces[0], swapped_face_frame)
                out.write(swapped_face_frame)
            else:
                # If no face detected, just write the original frame
                out.write(frame)

            # Increment frame count
            frame_count += 1

            # Calculate and print progress every 5%
            progress = (frame_count / total_frames) * 100
            if progress >= 5 * (frame_count // (total_frames // 20)):
                print(f"Processing: {int(progress)}% completed")

        cap.release()
        out.release()
        print(f"Video processing complete. Output saved to {dst_mp4}")


# Usage example:

# Initialize the FaceProcessor (use 'cuda' for GPU)
face_processor = FaceProcessor(device='cpu')

# Process the video by swapping faces from the target image
# face_processor.process_video('g1.mp4', 'from-face/girl-1.png', './g1-demo_output.mp4')
# face_processor.process_video('g1.mp4', 'from-face/girl-2.png', './g2-demo_output.mp4')
face_processor.process_video('g3.mp4', 'lilizhen.jpg', './lilizhen_output.mp4')
face_processor.process_video('g3.mp4', 'from-face/girl-1.png', './g1-demo_output.mp4')
face_processor.process_video('g3.mp4', 'from-face/girl-2.png', './g2-demo_output.mp4')