"""
Hugging Face Space — Face Swap with InsightFace inswapper_128.
Gradio app: webcam or upload a source face, swap it onto a target image.
"""

import os
import cv2
import numpy as np
import gradio as gr
import insightface
from insightface.app import FaceAnalysis
from huggingface_hub import hf_hub_download

HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_TARGET = os.path.join(HERE, "reference.png")

print("[boot] loading FaceAnalysis (buffalo_l)…", flush=True)
face_app = FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"])
face_app.prepare(ctx_id=0, det_size=(640, 640))

print("[boot] downloading inswapper_128…", flush=True)
MODEL_PATH = hf_hub_download(
    repo_id="ezioruan/inswapper_128.onnx",
    filename="inswapper_128.onnx",
)

print("[boot] loading inswapper_128…", flush=True)
swapper = insightface.model_zoo.get_model(
    MODEL_PATH, providers=["CPUExecutionProvider"]
)

print("[boot] ready", flush=True)


def largest_face(faces):
    return max(
        faces,
        key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]),
    )


def swap(source_img, target_img):
    if source_img is None:
        return None, "Please provide a source face (webcam or upload)."

    # Gradio gives RGB numpy — convert to BGR for OpenCV/InsightFace
    source_bgr = cv2.cvtColor(source_img, cv2.COLOR_RGB2BGR)

    if target_img is None:
        target_bgr = cv2.imread(DEFAULT_TARGET)
        if target_bgr is None:
            return None, "Default target image is missing."
    else:
        target_bgr = cv2.cvtColor(target_img, cv2.COLOR_RGB2BGR)

    source_faces = face_app.get(source_bgr)
    if not source_faces:
        return None, "No face detected in source image."
    src_face = largest_face(source_faces)

    target_faces = face_app.get(target_bgr)
    if not target_faces:
        return None, "No face detected in target image."
    tgt_face = largest_face(target_faces)

    result = swapper.get(target_bgr, tgt_face, src_face, paste_back=True)
    return cv2.cvtColor(result, cv2.COLOR_BGR2RGB), "Done."


with gr.Blocks(title="Face Swap — InsightFace") as demo:
    gr.Markdown(
        "# Face Swap — InsightFace `inswapper_128`\n"
        "Take a webcam photo (or upload), click **Swap**. "
        "Leave the target empty to use the default Denmark fan, "
        "or upload your own target image."
    )
    with gr.Row():
        with gr.Column():
            src_in = gr.Image(
                sources=["webcam", "upload"],
                type="numpy",
                label="Your face (webcam or upload)",
            )
            tgt_in = gr.Image(
                sources=["upload"],
                type="numpy",
                label="Target image (optional — defaults to reference.png)",
                value=DEFAULT_TARGET,
            )
            btn = gr.Button("Swap", variant="primary")
        with gr.Column():
            out_img = gr.Image(label="Result", type="numpy")
            status = gr.Textbox(label="Status", interactive=False)

    btn.click(swap, inputs=[src_in, tgt_in], outputs=[out_img, status])

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