# TÊN TỆP: app.py (Mã đã được dọn dẹp và sử dụng cache)

import gradio as gr
import cv2 
import numpy as np
import insightface
from insightface.app import FaceAnalysis
import os
import urllib.request
import ssl

# --- SỬA LỖI: CÀI ĐẶT ĐƯỜNG DẪN MÔ HÌNH THỦ CÔNG ---
# Thiết lập ngữ cảnh SSL cho tải xuống
ssl._create_default_https_context = ssl._create_unverified_context

MODEL_NAME = 'inswapper_128.onnx'
# URL đã được xác nhận hoạt động trên Hugging Face
MODEL_URL = 'https://huggingface.co/xingren23/comfyflow-models/resolve/976de8449674de379b02c144d0b3cfa2b61482f2/insightface/inswapper_128.onnx' 

# Thư mục cache mặc định của InsightFace
cache_dir = os.path.join(os.path.expanduser('~'), '.insightface', 'models')
os.makedirs(cache_dir, exist_ok=True)
MODEL_PATH = os.path.join(cache_dir, MODEL_NAME)


# Hàm Tải Mô hình Bằng Tay (để đảm bảo tệp tồn tại)
def download_model(url, path):
    if not os.path.exists(path):
        print(f"Downloading {os.path.basename(path)} from {url} to {path}...")
        try:
            # Tải tệp vào thư mục cache chính xác
            urllib.request.urlretrieve(url, path)
            print(f"Successfully downloaded {os.path.basename(path)}.")
        except Exception as e:
            print(f"FATAL ERROR: Failed to download model {os.path.basename(path)}: {e}")
            raise e
    return path

# TẢI MÔ HÌNH VÀO THƯ MỤC INSIGHTFACE CHUẨN
local_model_path = download_model(MODEL_URL, MODEL_PATH)

# --- KHỞI TẠO ỨNG DỤNG ---
try:
    # Tải mô hình Face Analysis
    app = FaceAnalysis(
        name='buffalo_l', 
        providers=['CPUExecutionProvider']
    )
    app.prepare(ctx_id=0, det_size=(640, 640)) 

    # Tải mô hình Face Swapper bằng đường dẫn cục bộ đã tải
    # SỬA LỖI: Chỉ định đường dẫn mô hình (local_model_path)
    swapper = insightface.model_zoo.get_model(local_model_path, providers=['CPUExecutionProvider'])
    
except Exception as e:
    print(f"FATAL ERROR during model initialization: {e}")
    raise e


# HÀM HOÁN ĐỔI KHUÔN MẶT THỰC TẾ
def face_swap_function(source_path, target_path):
    # ... (Giữ nguyên logic hàm)
    try:
        img_source = cv2.imread(source_path)
        img_target = cv2.imread(target_path)
    except Exception:
        return cv2.cvtColor(np.zeros((200, 200, 3), dtype=np.uint8), cv2.COLOR_BGR2RGB)

    if img_source is None or img_target is None:
        return cv2.cvtColor(np.zeros((200, 200, 3), dtype=np.uint8), cv2.COLOR_BGR2RGB)

    faces_source = app.get(img_source)
    if not faces_source:
        return cv2.cvtColor(img_target, cv2.COLOR_BGR2RGB)
        
    source_face = faces_source[0]
    
    faces_target = app.get(img_target)
    if not faces_target:
        return cv2.cvtColor(img_target, cv2.COLOR_BGR2RGB)
        
    target_face = faces_target[0]
    
    result_img = swapper.get(img_target, target_face, source_face, paste_back=True)
    
    return cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)

# TẠO GIAO DIỆN GRADIO
demo = gr.Interface(
    fn=face_swap_function,
    inputs=[
        gr.Image(type="filepath", label="Ảnh Nguồn (Khuôn mặt bạn muốn dùng)"),
        gr.Image(type="filepath", label="Ảnh Đích (Khuôn mặt bạn muốn thay thế)")
    ],
    outputs=gr.Image(type="numpy", label="Kết quả Hoán đổi Khuôn mặt (Face Swap)"),
    title="Ứng dụng Hoán đổi Khuôn mặt INSIGHTFACE Hoàn chỉnh"
)

# Chạy ứng dụng
demo.launch()