import gradio as gr
import subprocess
import os
import shutil
import tempfile
from datetime import datetime
from huggingface_hub import (
    login,
    hf_hub_download,
    upload_file,
    upload_folder,
    HfApi,
)

# -------------------------------------------------
# Cấu hình (đặt các biến môi trường trước khi chạy)
# -------------------------------------------------
REPO_ID = os.getenv("HFPATH")          # ví dụ: "username/dataset-name"
HF_TOKEN = os.getenv("MAGIC")          # token Hugging Face của bạn
login(HF_TOKEN)

api = HfApi()


# -------------------------------------------------
# Hàm phụ: tải xuống tất cả file trong một subfolder
# -------------------------------------------------
def download_from_hf(subfolder: str):
    """Trả về danh sách đường dẫn cục bộ của các file trong `subfolder`."""
    downloaded_files = []

    # Lấy danh sách mọi file trong repo dataset
    all_files = api.list_repo_files(repo_id=REPO_ID, repo_type="dataset")

    # Chỉ giữ lại những file bắt đầu bằng subfolder yêu cầu
    for file in all_files:
        if file.startswith(f"{subfolder}/"):
            local_path = hf_hub_download(
                repo_id=REPO_ID,
                filename=file,
                repo_type="dataset",
            )
            downloaded_files.append(local_path)

    return downloaded_files


# -------------------------------------------------
# Hàm phụ: tải lên một file đơn lẻ vào thư mục theo ngày
# -------------------------------------------------
def upload_to_hf(filepath: str):
    filename = os.path.basename(filepath)
    subfolder = datetime.now().strftime("%Y%m%d")
    path_in_repo = f"{subfolder}/{filename}"

    upload_file(
        path_or_fileobj=filepath,
        path_in_repo=api,
        repo_id=REPO_ID,
        repo_type="dataset",
    )
    return f"https://huggingface.co/datasets/{REPO_ID}/blob/main/{path_in_repo}"


# -------------------------------------------------
# Hàm phụ: tải lên toàn bộ thư mục vào thư mục theo ngày
# -------------------------------------------------
def upload_folder_to_hf(folder_path: str):
    subfolder = datetime.now().strftime("%Y%m%d")
    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}"


# -------------------------------------------------
# Hàm chính được Gradio gọi
# -------------------------------------------------
def run_scripts(subfolder: str, source):
    """
    1. Tải về các file mục tiêu từ `subfolder`.
    2. Chạy `run.py` cho mỗi file, tạo file output.
    3. Sao chép các output vào thư mục tạm.
    4. Tải lên thư mục tạm lên Hugging Face.
    Trả về danh sách tên file output.
    """
    output_files = []
    target_files = download_from_hf(subfolder)

    # Thư mục tạm để gom kết quả
    with tempfile.TemporaryDirectory() as temp_dir:
        print(f"Thư mục tạm được tạo: {temp_dir}")

        for target_file in target_files:
            # Giữ nguyên phần mở rộng của file gốc
            _, ext = os.path.splitext(target_file)
            timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
            output_path = f"output_{timestamp}{ext}"

            # Lệnh gọi script swap mặt
            cmd = [
                "python3",
                "run.py",
                "-s",
                source.name,
                "-t",
                target_file,
                "-o",
                output_path,
                "--frame-processor",
                "face_swapper",
                "--many-faces",
            ]

            # Thực thi, sẽ ném lỗi nếu có vấn đề
            subprocess.run(cmd, check=True)

            output_files.append(output_path)
            print(f"Tạo: {output_path}")

            # Sao chép vào thư mục tạm
            dest_path = os.path.join(temp_dir, os.path.basename(output_path))
            shutil.copy(output_path, dest_path)
            print(f"Đã sao chép {output_path} → {dest_path}")

        # Tải lên toàn bộ thư mục tạm
        upload_folder_to_hf(temp_dir)

    return output_files


# -------------------------------------------------
# Giao diện Gradio
# -------------------------------------------------
iface = gr.Interface(
    fn=run_scripts,
    inputs=[
        "text",   # tên subfolder
        "file",   # ảnh nguồn
    ],
    outputs="files",
    title="Face Swapper",
    description=(
        "Nhập tên subfolder trong dataset và tải lên ảnh nguồn để thực hiện "
        "đổi mặt."
    ),
    live=False,
)

iface.launch()