import numpy as np
from PIL import Image
import insightface
from insightface.app import FaceAnalysis

import os
import requests

MODEL_URL = "https://huggingface.co/kaizma/face-swap-inswapper/resolve/da20be1c8ba9b074d52c6a0540f8935d3e3605e5/inswapper_128.onnx"
MODEL_PATH = "inswapper_128.onnx"

# Download model if not exists
if not os.path.exists(MODEL_PATH):
    print("Downloading model...")
    r = requests.get(MODEL_URL, stream=True)
    with open(MODEL_PATH, "wb") as f:
        for chunk in r.iter_content(chunk_size=8192):
            f.write(chunk)
    print("Model downloaded.")


def swap_faces(src_pil, tgt_pil):
    # Convert PIL to numpy (BGR)
    src = np.array(src_pil)[:, :, ::-1].copy()
    tgt = np.array(tgt_pil)[:, :, ::-1].copy()

    # Initialize face analysis and swapper
    app = FaceAnalysis(name='buffalo_l')
    app.prepare(ctx_id=0, det_size=(640, 640))
    model_path = r"C:\Users\Dell\Desktop\Face Swap\inswapper_128.onnx"
    swapper = insightface.model_zoo.get_model(model_path)

    # Detect faces
    src_faces = app.get(src)
    tgt_faces = app.get(tgt)
    if len(src_faces) == 0 or len(tgt_faces) == 0:
        raise ValueError("No face detected in one of the images.")

    source_face = src_faces[0]
    res = tgt.copy()
    for face in tgt_faces:
        res = swapper.get(res, face, source_face, paste_back=True)

    # Convert back to PIL (RGB)
    res_pil = Image.fromarray(res[:, :, ::-1])
    return res_pil