import gradio as gr
import insightface
from insightface.app import FaceAnalysis
import numpy as np
from PIL import Image
import cv2
import tempfile
import os

MODEL_PATH = "inswapper_128.onnx"

# --------- Load models once ---------
print("Loading models...")
face_app = FaceAnalysis(name="buffalo_l")
face_app.prepare(ctx_id=0, det_size=(640, 640))  # ctx_id=0 -> CPU on Spaces free tier
swapper = insightface.model_zoo.get_model(MODEL_PATH)
print("Models loaded.")

# --------- IMAGE SWAP ---------
def swap_faces_image(source, target):
    if source is None or target is None:
        raise gr.Error("Please upload both source and target images.")

    # Gradio gives numpy arrays (H, W, 3) in RGB
    src = np.array(source)
    dst = np.array(target)

    src_faces = face_app.get(src)
    dst_faces = face_app.get(dst)

    if len(src_faces) == 0:
        raise gr.Error("No face found in source image.")
    if len(dst_faces) == 0:
        raise gr.Error("No face found in target image.")

    src_face = src_faces[0]
    dst_face = dst_faces[0]

    result = swapper.get(dst.copy(), dst_face, src_face, paste_back=True)
    return Image.fromarray(result)

# --------- VIDEO SWAP ---------
def swap_faces_video(source_image, video_file):
    if source_image is None:
        raise gr.Error("Please upload a source face image.")
    if video_file is None:
        raise gr.Error("Please upload a target video (mp4).")

    # source_image: numpy RGB
    src = np.array(source_image)
    src_faces = face_app.get(src)
    if len(src_faces) == 0:
        raise gr.Error("No face found in source image.")
    src_face = src_faces[0]

    # video_file is a file-like object; get its path
    video_path = video_file.name

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise gr.Error("Could not open uploaded video.")

    fps = cap.get(cv2.CAP_PROP_FPS) or 25
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # temp output file
    tmp_out = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    tmp_out_path = tmp_out.name
    tmp_out.close()

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(tmp_out_path, fourcc, fps, (w, h))

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    print(f"[INFO] Total frames: {total_frames}")

    frame_idx = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame_idx += 1
        if frame_idx % 10 == 0:
            print(f"[INFO] Frame {frame_idx}/{total_frames}")

        # BGR -> RGB
        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        dst_faces = face_app.get(rgb)

        if len(dst_faces) > 0:
            dst_face = dst_faces[0]
            swapped_rgb = swapper.get(rgb.copy(), dst_face, src_face, paste_back=True)
            out_frame = cv2.cvtColor(swapped_rgb, cv2.COLOR_RGB2BGR)
            writer.write(out_frame)
        else:
            writer.write(frame)

    cap.release()
    writer.release()

    print(f"[INFO] Video finished: {tmp_out_path}")
    return tmp_out_path

# --------- Gradio UI (Image + Video) ---------
with gr.Blocks() as demo:
    gr.Markdown("## InsightFace FaceSwap (Image & Video) on HuggingFace")

    with gr.Tab("Image Swap"):
        src_img = gr.Image(type="numpy", label="Your Face (Source)")
        tgt_img = gr.Image(type="numpy", label="Target Face")
        out_img = gr.Image(label="Swapped Result")

        btn_img = gr.Button("Swap Image")
        btn_img.click(fn=swap_faces_image, inputs=[src_img, tgt_img], outputs=out_img)

    with gr.Tab("Video Swap"):
        src_vid_img = gr.Image(type="numpy", label="Your Face (Source)")
        tgt_vid = gr.File(label="Target Video (mp4)")
        out_vid = gr.Video(label="Swapped Video")

        btn_vid = gr.Button("Swap Video")
        btn_vid.click(fn=swap_faces_video, inputs=[src_vid_img, tgt_vid], outputs=out_vid)

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