import gradio as gr
import os
import sys

# roop folder ko system path mein add karna
sys.path.append('roop')

# Sahi jagah se function import karna
from roop.core import run as roop_run

# --- Gradio Function ---
def face_swap(source_img, target_media):
    if source_img is None or target_media is None:
        raise gr.Error("Please upload both source image and target media.")

    # ===== YAHAN MERI GALTI THI, AB THEEK HAI =====
    # OLD: source_path = source_img.name
    # NEW: Gradio 'filepath' type mein seedha path (string) deta hai.
    source_path = source_img
    target_path = target_media
    # ===============================================
    
    # Output file ka naam tay karna
    output_filename = os.path.basename(target_path)
    os.makedirs("output", exist_ok=True)
    output_path = os.path.join("output", f"result_{output_filename}")

    # Roop ko chalane ke liye arguments set karna
    args = [
        "run.py",
        "--source", source_path,
        "--target", target_path,
        "--output", output_path,
        "--execution-provider", "cpu",
        "-e"
    ]
    sys.argv = args
    
    print(f"Running roop with args: {sys.argv}")

    # Asli face swap process yahan run hoga
    try:
        roop_run()
    except Exception as e:
        # Error aane par user ko message dikhana
        raise gr.Error(f"Face swap failed. Error: {str(e)}")

    print(f"Processing complete. Result at: {output_path}")

    # Result ko UI mein dikhana
    return output_path

# --- Gradio UI ---
with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🎭 Roop Face Swap AI")
    gr.Markdown("Upload a source face and a target image/video, then click 'Swap Face'.")
    
    with gr.Row():
        source_image = gr.Image(label="Source Face (Single Face)", type="filepath")
        target_media = gr.File(label="Target Image or Video", type="filepath")
    
    submit_btn = gr.Button("Swap Face", variant="primary")
    result_output = gr.File(label="Result", interactive=False)
    
    submit_btn.click(
        fn=face_swap,
        inputs=[source_image, target_media],
        outputs=result_output
    )
    
    gr.Markdown("---")
    gr.Markdown("Created with help from AI. Make sure to use this tool responsibly.")

app.launch()