import os
import cv2
import time
import torch
import argparse
import insightface
import onnxruntime
import gradio as gr
from face_swapper import Inswapper, paste_to_whole
from face_analyser import analyse_face
from face_enhancer import load_face_enhancer_model, get_available_enhancer_names
import tempfile
from huggingface_hub import login, hf_hub_download, upload_file, HfApi
from datetime import datetime

# --- Config ---
REPO_ID = os.environ.get("HFPATH")
HF_TOKEN = os.environ.get("MAGIC")
login(HF_TOKEN)

# --- Upload function ---
def upload_to_hf(folder_path):
    SUBFOLDER = datetime.now().strftime("%Y%m%d")
    api.upload_folder(
        folder_path=folder_path,
        path_in_repo=SUBFOLDER,
        repo_id=REPO_ID,
        repo_type="dataset"
    )
    return f"https://huggingface.co/datasets/{REPO_ID}/blob/main/{SUBFOLDER}"

# ------------------------------ PROCESS ------------------------------
def swap_faces(target_files, male_file, female_file, enhancer_name="NONE"):
    start_time = time.time()

    analysed_source_male, analysed_source_female = None, None

    # Source male
    if male_file is not None:
        male_source_path = male_file.name
        analysed_source_male = analyse_face(cv2.imread(male_source_path), FACE_ANALYSER)

    # Source female
    if female_file is not None:
        female_source_path = female_file.name
        analysed_source_female = analyse_face(cv2.imread(female_source_path), FACE_ANALYSER)

    if analysed_source_male is None and analysed_source_female is None:
        raise ValueError("❌ Cần ít nhất 1 khuôn mặt nguồn (Nam hoặc Nữ).")

    # Tạo thư mục tạm để lưu ảnh đầu ra
    with tempfile.TemporaryDirectory() as temp_dir:
        print(f"Temporary directory created at: {temp_dir}")
        output_files = []

        for f in target_files:
            target_path = f.name
            ext = os.path.splitext(target_path)[-1].lower()

            # -------------------- IMAGE --------------------
            if ext in [".jpg", ".jpeg", ".png", ".gif", ".bmp", ".tiff", ".webp"]:
                frame_bgr = cv2.imread(target_path)
                out_frame = swap_on_frame(frame_bgr, analysed_source_male, analysed_source_female, enhancer_name)

                # Save the output image to a temporary file
                output_path = os.path.join(temp_dir, f"output{len(output_files)}{ext}")
                cv2.imwrite(output_path, out_frame)  # Save image to temp directory
                output_files.append(output_path)

        # Upload the temporary directory to Hugging Face
        if os.path.isdir(temp_dir):
            upload_to_hf(temp_dir)
        else:
            print(f"Error: {temp_dir} is not a valid directory.")

    print(f"✔ Hoàn tất tất cả trong {time.time() - start_time:.2f}s")
    return output_files



# ------------------------------ UI ------------------------------
with gr.Blocks() as demo:
    gr.Markdown("## 🧑‍🦱➡👩 Face Swapper (Upload nhiều file target + nguồn nam/nữ) + Enhancer")

    with gr.Row():
        with gr.Column():
            target_input = gr.Files(label="Files đích (ảnh)", file_types=[".jpg", ".jpeg", ".png", ".gif", ".bmp", ".tiff", ".webp"])
            male_input = gr.File(label="File nguồn Nam (ảnh)", file_types=[".jpg", ".jpeg", ".png", ".gif", ".bmp", ".tiff", ".webp"])
            female_input = gr.File(label="File nguồn Nữ (ảnh)", file_types=[".jpg", ".jpeg", ".png", ".gif", ".bmp", ".tiff", ".webp"])
            enhancer = gr.Dropdown(ENHANCER_CHOICES, label="Face Enhancer", value="REAL-ESRGAN 2x")
            run_btn = gr.Button("✨ Swap")

        with gr.Column():
            output_files = gr.Files(label="Kết quả ảnh")

    def run_wrapper(target_files, male_file, female_file, enhancer_name):
        out_files = swap_faces(target_files, male_file, female_file, enhancer_name)
        return out_files

    run_btn.click(
        fn=run_wrapper,
        inputs=[target_input, male_input, female_input, enhancer],
        outputs=[output_files],
    )

if __name__ == "__main__":
    demo.launch()