import os
os.environ["OMP_NUM_THREADS"] = "1"
import gradio as gr
import cv2
import shutil
import uuid
import insightface
from insightface.app import FaceAnalysis
from huggingface_hub import hf_hub_download
import subprocess


# -------------------------------------------------
# Paths
# -------------------------------------------------
REPO_ID = "HariLogicgo/face_swap_models"  # <- your HF repo for models
BASE_DIR = "./workspace"
UPLOAD_DIR = os.path.join(BASE_DIR, "uploads")
RESULT_DIR = os.path.join(BASE_DIR, "results")
MODELS_DIR = "./models"

os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(RESULT_DIR, exist_ok=True)

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

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
    )

# -------------------------------------------------
# Initialize face analysis and 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'])

# -------------------------------------------------
# CodeFormer setup
# -------------------------------------------------
CODEFORMER_PATH = "CodeFormer/inference_codeformer.py"

def ensure_codeformer():
    if not os.path.exists("CodeFormer"):
        subprocess.run("git clone https://github.com/sczhou/CodeFormer.git", shell=True)
        subprocess.run("pip install -r CodeFormer/requirements.txt", shell=True)
        subprocess.run("python CodeFormer/basicsr/setup.py develop", shell=True)
        subprocess.run("python CodeFormer/scripts/download_pretrained_models.py facelib", shell=True)
        subprocess.run("python CodeFormer/scripts/download_pretrained_models.py CodeFormer", shell=True)

ensure_codeformer()

# -------------------------------------------------
# Pipeline Function
# -------------------------------------------------
def face_swap_and_enhance(src_img, tgt_img, fidelity=0.7, background_enhance=True, face_upsample=True):
    try:
        src_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
        tgt_bgr = cv2.cvtColor(tgt_img, cv2.COLOR_RGB2BGR)

        src_faces = app.get(src_bgr)
        tgt_faces = app.get(tgt_bgr)
        if not src_faces or not tgt_faces:
            return None, None, "❌ Face not detected in one of the images."

        shutil.rmtree(UPLOAD_DIR, ignore_errors=True)
        shutil.rmtree(RESULT_DIR, ignore_errors=True)
        os.makedirs(UPLOAD_DIR, exist_ok=True)
        os.makedirs(RESULT_DIR, exist_ok=True)

        unique_name = f"swapped_{uuid.uuid4().hex[:8]}.jpg"
        swapped_path = os.path.join(UPLOAD_DIR, unique_name)
        swapped_bgr = swapper.get(tgt_bgr, tgt_faces[0], src_faces[0])
        cv2.imwrite(swapped_path, swapped_bgr)

        cmd = f"python {CODEFORMER_PATH} -w {fidelity:.2f} --input_path {UPLOAD_DIR} --output_path {RESULT_DIR}"
        if background_enhance:
            cmd += " --bg_upsampler realesrgan"
        if face_upsample:
            cmd += " --face_upsample"

        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        if result.returncode != 0:
            return None, None, f"❌ CodeFormer failed:\n{result.stderr}"

        final_path = None
        for root, _, files in os.walk(RESULT_DIR):
            for f in files:
                if f.endswith((".png", ".jpg")):
                    final_path = os.path.join(root, f)
                    break
            if final_path:
                break

        if not final_path or not os.path.exists(final_path):
            return None, None, "❌ CodeFormer output missing."

        final_img = cv2.cvtColor(cv2.imread(final_path), cv2.COLOR_BGR2RGB)
        return final_img, final_path, ""

    except Exception as e:
        return None, None, f"❌ Error: {str(e)}"

# -------------------------------------------------
# Gradio Interface
# -------------------------------------------------
with gr.Blocks() as demo:
    gr.Markdown("## 🧑‍🤝‍🧑 Face Swap + CodeFormer Enhancement")
    with gr.Row():
        src_input = gr.Image(type="numpy", label="Upload Source Face")
        tgt_input = gr.Image(type="numpy", label="Upload Target Image")
    with gr.Row():
        fidelity = gr.Slider(0, 1, value=0.7, step=0.01, label="CodeFormer Fidelity")
        bg = gr.Checkbox(value=True, label="Enhance Background")
        face_up = gr.Checkbox(value=True, label="Face Upsample")
    btn = gr.Button("🚀 Run Face Swap + Enhance")

    output_img = gr.Image(type="numpy", label="Enhanced Output")
    download = gr.File(label="⬇️ Download Enhanced Image")
    error_box = gr.Textbox(label="Logs / Errors", interactive=False)

    def process(src, tgt, f, b, fu):
        img, path, err = face_swap_and_enhance(src, tgt, f, b, fu)
        return img, path, err

    btn.click(process, [src_input, tgt_input, fidelity, bg, face_up],
              [output_img, download, error_box])

demo.launch()