import gradio as gr
import subprocess
import os
from datetime import datetime
from PIL import Image
from huggingface_hub import login, upload_file

# --- Config ---
REPO_ID = os.environ.get("HFPATH")
HF_TOKEN = os.environ.get("MAGIC")
login(HF_TOKEN)
def upload_to_hf(filepath):
    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=path_in_repo,
        repo_id=REPO_ID,
        repo_type="dataset"
    )
    return f"https://huggingface.co/datasets/{REPO_ID}/blob/main/{filename}"

def run_scripts(target, source, use_face_enhancer):
    output_images = []  # Danh sách để lưu các ảnh đã xử lý
    for target_file in target:
        target_extension = os.path.splitext(target_file.name)[-1]  # Lấy phần mở rộng của tệp
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        output_path1 = f"output_{timestamp}{target_extension}"  # Đường dẫn tệp đầu ra

        # Chạy lệnh subprocess để xử lý ảnh
        cmd1 = [
            "python3", "run.py", "-s", source.name, "-t", target_file.name,
            "-o", output_path1, "--frame-processor", "face_swapper"
        ]
        subprocess.run(cmd1)

        # Mở ảnh đã được xử lý và thêm vào danh sách
        output_image = Image.open(output_path1)
        upload_to_hf(output_path1)
        output_images.append(output_image)

    return output_images  # Trả về danh sách ảnh đã xử lý

# Giao diện Gradio
iface = gr.Interface(
    fn=run_scripts,
    inputs=[
        gr.Files(label="Target Files"),  # Đầu vào là tệp tin mục tiêu (ảnh/video)
        gr.File(label="Source File"),    # Đầu vào là tệp nguồn (ảnh)
        gr.Checkbox(label="Use Face Enhancer")  # Tùy chọn cho face enhancer
    ],
    outputs=gr.Gallery(label="Output Images"),  # Đầu ra là các ảnh đã được xử lý
    title="Face Swapper",
    description="Upload a target image/video and a source image to swap faces.",
    live=False
)

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