import os
import gradio as gr
import subprocess
from lib.config import args

def process_video(video_path, task):
    """
    Function to process the input video and run the NeuralBody pipeline.
    """
    # Save uploaded video locally
    input_video = "input_video.mp4"
    os.system(f"cp {video_path} {input_video}")

    # Map tasks to functions in run.py
    task_map = {
        "Dataset Processing": "dataset",
        "Network Inference": "network",
        "Evaluation": "evaluate",
        "Visualization": "visualize",
    }

    if task not in task_map:
        return "Invalid task selected!"

    # Run corresponding function in run.py
    args.type = task_map[task]  # Set the correct function call in run.py
    subprocess.run(["python", "run.py"], check=True)

    return f"Task '{task}' completed! Check the output directory."

# Gradio UI
iface = gr.Interface(
    fn=process_video,
    inputs=[
        gr.Video(label="Upload Video"),
        gr.Radio(["Dataset Processing", "Network Inference", "Evaluation", "Visualization"], label="Select Task"),
    ],
    outputs="text",
    title="NeuralBody: Video-Based 3D Reconstruction",
    description="Upload a video and choose a task to perform using the NeuralBody pipeline."
)

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