import cv2
import insightface
from insightface.app import FaceAnalysis
import gradio as gr
from huggingface_hub import hf_hub_download
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
import os

# -------------------------------------------------
# Hugging Face repo ID where models are stored
# -------------------------------------------------
REPO_ID = "HariLogicgo/face_swap_models"
MODELS_DIR = "./models"

# -------------------------------------------------
# Download models from HF Hub
# -------------------------------------------------
# InSwapper model
inswapper_path = hf_hub_download(
    repo_id=REPO_ID,
    filename="models/inswapper_128.onnx",
    repo_type="model",
    local_dir=MODELS_DIR
)

# Buffalo_l ONNX files
buffalo_files = [
    "1k3d68.onnx",
    "2d106det.onnx",
    "genderage.onnx",
    "det_10g.onnx",
    "w600k_r50.onnx"
]

for f in buffalo_files:
    hf_hub_download(
        repo_id=REPO_ID,
        filename=f"models/buffalo_l/{f}",
        repo_type="model",
        local_dir=MODELS_DIR
    )

# RealESRGAN model
realesrgan_path = hf_hub_download(
    repo_id=REPO_ID,
    filename="models/RealESRGAN_x4plus.pth",
    repo_type="model",
    local_dir=MODELS_DIR
)

# -------------------------------------------------
# Initialize Face Analysis + Swapper
# -------------------------------------------------
app = FaceAnalysis(name="buffalo_l", root=MODELS_DIR, providers=['CPUExecutionProvider'])
app.prepare(ctx_id=0, det_size=(640, 640))

swapper = insightface.model_zoo.get_model(inswapper_path, providers=['CPUExecutionProvider'])

# -------------------------------------------------
# Initialize RealESRGAN Upscaler (CPU mode)
# -------------------------------------------------
model = RRDBNet(
    num_in_ch=3,
    num_out_ch=3,
    num_feat=64,
    num_block=23,
    num_grow_ch=32,
    scale=4
)

upscaler = RealESRGANer(
    scale=4,
    model_path=realesrgan_path,
    dni_weight=None,
    model=model,
    tile=0,
    tile_pad=10,
    pre_pad=0,
    half=False,    # must be False for CPU
    gpu_id=None    # ensures CPU mode
)

# -------------------------------------------------
# Face swap + upscale function
# -------------------------------------------------
def swap_face_image(source_img, target_img):
    src_img = cv2.cvtColor(source_img, cv2.COLOR_RGB2BGR)
    tgt_img = cv2.cvtColor(target_img, cv2.COLOR_RGB2BGR)

    src_faces = app.get(src_img)
    tgt_faces = app.get(tgt_img)

    if len(src_faces) == 0 or len(tgt_faces) == 0:
        return None  # No face detected

    swapped = swapper.get(tgt_img, tgt_faces[0], src_faces[0])

    # Upscale with RealESRGAN
    upscaled, _ = upscaler.enhance(swapped, outscale=4)

    # Convert back to RGB for Gradio
    result = cv2.cvtColor(upscaled, cv2.COLOR_BGR2RGB)
    return result

# -------------------------------------------------
# Gradio Interface
# -------------------------------------------------
with gr.Blocks() as demo:
    gr.Markdown("### Face Swap + Upscale (CPU, HuggingFace Models)")

    with gr.Row():
        src_img = gr.Image(type="numpy", label="Source Image")
        tgt_img = gr.Image(type="numpy", label="Target Image")

    output_img = gr.Image(type="numpy", label="Swapped & Upscaled Output")
    btn1 = gr.Button("Swap + Upscale")

    btn1.click(swap_face_image, inputs=[src_img, tgt_img], outputs=output_img)

# -------------------------------------------------
# Launch
# -------------------------------------------------
demo.launch()