import gradio as gr
import subprocess
import os
import shutil
from datetime import datetime
from PIL import Image
import uuid

OUTPUT_DIR = "outputs"
os.makedirs(OUTPUT_DIR, exist_ok=True)


def run_scripts(target_files, source_file, use_face_enhancer):

    if not target_files or not source_file:
        return []

    output_results = []

    for target_file in target_files:

        # ===== Unique filename =====
        unique_id = uuid.uuid4().hex[:8]
        extension = os.path.splitext(target_file.name)[-1]
        output_path = os.path.join(
            OUTPUT_DIR,
            f"swap_{unique_id}{extension}"
        )

        # ===== Build command =====
        cmd = [
            "python3",
            "run.py",
            "-s", source_file.name,
            "-t", target_file.name,
            "-o", output_path,
            "--frame-processor", "face_swapper"
        ]

        if use_face_enhancer:
            cmd += ["--face-enhancer"]

        print("Running:", " ".join(cmd))

        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True
            )

            if result.returncode != 0:
                print("Error:")
                print(result.stderr)
                continue

        except Exception as e:
            print("Subprocess crashed:", str(e))
            continue

        # ===== Check output exists =====
        if not os.path.exists(output_path):
            print("Output not found:", output_path)
            continue

        # ===== Handle file types =====
        ext = extension.lower()

        try:
            if ext in [".png", ".jpg", ".jpeg", ".webp"]:
                img = Image.open(output_path).convert("RGB")
                output_results.append(img)

            elif ext == ".gif":
                # Return gif path directly (Gallery hiển thị được)
                output_results.append(output_path)

            elif ext in [".mp4", ".mov", ".avi"]:
                output_results.append(output_path)

            else:
                output_results.append(output_path)

        except Exception as e:
            print("Error loading result:", str(e))
            continue

    return output_results


iface = gr.Interface(
    fn=run_scripts,
    inputs=[
        gr.Files(label="Target Files (Image / GIF / Video)"),
        gr.File(label="Source Face Image"),
        gr.Checkbox(label="Use Face Enhancer", value=False)
    ],
    outputs=gr.Gallery(label="Results"),
    title="Face Swapper",
    description="Upload target images/videos and a source image to swap faces.",
    live=False
)

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