import gradio as gr
import os
import sys

sys.path.append('roop')
from roop.core import run as roop_run

def face_swap_logic(source_img, target_media):
    if source_img is None or target_media is None:
        raise gr.Error("Dono file upload karein.")

    source_path = source_img
    target_path = target_media
    
    # Check karein ki target video hai ya image
    is_video = target_path.lower().endswith(('.mp4', '.mov', '.avi'))
    
    output_filename = os.path.basename(target_path)
    os.makedirs("output", exist_ok=True)
    output_path = os.path.join("output", f"result_{output_filename}")

    args = [ "run.py", "--source", source_path, "--target", target_path, "--output", output_path, "--execution-provider", "cpu" ]
    sys.argv = args
    
    try:
        roop_run()
    except SystemExit: pass
    except Exception as e:
        raise gr.Error(f"Face swap fail hua. Roop ke andar error: {str(e)}")

    # Ab hum behtar tarike se check karenge
    if not os.path.exists(output_path):
        if is_video:
            # Video ke liye alag error message
            raise gr.Error("Video processing fail hua. Free server par time ya memory (RAM) kam pad gayi. Please chota video (5-10 sec) try karein, ya GPU server use karein.")
        else:
            # Image ke liye alag error message
            raise gr.Error("Image processing fail hua. Aisa tab hota hai jab image mein chehra nahi milta.")

    return output_path

# --- UI Code (Bilkul Same, Koi Change Nahi) ---
with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🎭 Roop Face Swap AI")
    with gr.Tabs():
        with gr.TabItem("🖼️ Swap on Image"):
            with gr.Row():
                source_face_img = gr.Image(label="Source Face", type="filepath")
                target_image = gr.Image(label="Target Image", type="filepath")
            submit_btn_image = gr.Button("Swap Face on Image")
            result_image = gr.Image(label="Result Image")

        with gr.TabItem("🎬 Swap on Video"):
            with gr.Row():
                source_face_vid = gr.Image(label="Source Face", type="filepath")
                target_video = gr.Video(label="Target Video")
            submit_btn_video = gr.Button("Swap Face on Video")
            result_video = gr.Video(label="Result Video")

    submit_btn_image.click(fn=face_swap_logic, inputs=[source_face_img, target_image], outputs=result_image)
    submit_btn_video.click(fn=face_swap_logic, inputs=[source_face_vid, target_video], outputs=result_video)

app.launch()