import os
import gc
import cv2
import gradio as gr
import numpy as np
from PIL import Image
from insightface.app import FaceAnalysis
from insightface.model_zoo import get_model

# -----------------------------
# OpenCV CPU control
# -----------------------------
cv2.setNumThreads(2)

# -----------------------------
# Model path (ROOT folder)
# -----------------------------
SWAP_MODEL_PATH = "inswapper_128.onnx"

if not os.path.exists(SWAP_MODEL_PATH):
    raise RuntimeError(
        "❌ inswapper_128.onnx not found.\n"
        "Upload it in the SAME folder as app.py"
    )

# -----------------------------
# Face detector (High Quality)
# -----------------------------
face_app = FaceAnalysis(
    providers=["CPUExecutionProvider"],
    allowed_modules=["detection", "recognition"]
)

face_app.prepare(
    ctx_id=0,
    det_size=(640, 640)  # 16 GB RAM safe
)

# -----------------------------
# Load face swap model (LOCAL)
# -----------------------------
swapper = get_model(
    SWAP_MODEL_PATH,
    providers=["CPUExecutionProvider"]
)

# -----------------------------
# Face swap function
# -----------------------------
def face_swap(source_img, target_img):
    try:
        if source_img is None or target_img is None:
            return None

        src = np.array(source_img.convert("RGB"))
        tgt = np.array(target_img.convert("RGB"))

        src_faces = face_app.get(src)
        tgt_faces = face_app.get(tgt)

        if len(src_faces) == 0 or len(tgt_faces) == 0:
            return None

        result = swapper.get(
            tgt,
            tgt_faces[0],
            src_faces[0],
            paste_back=True
        )

        return Image.fromarray(result)

    finally:
        del src, tgt, src_faces, tgt_faces
        gc.collect()

# -----------------------------
# Gradio Interface
# -----------------------------
demo = gr.Interface(
    fn=face_swap,
    inputs=[
        gr.Image(type="pil", label="Source Face"),
        gr.Image(type="pil", label="Target Image")
    ],
    outputs=gr.Image(type="pil"),
    analytics_enabled=False
)

# -----------------------------
# Launch (NO deprecated args)
# -----------------------------
demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    share=False
)