import os
import gradio as gr
import subprocess
import datetime
import sys

def run_command(command):
    """Run a shell command and print its output."""
    print(f"Running command: {' '.join(command)}")
    try:
        subprocess.check_call(command, shell=True)
    except subprocess.CalledProcessError as e:
        print(f"Error running command {command}: {e}")
        sys.exit(1)

def check_for_mp4_in_outputs(given_folder):
    # Define the path to the outputs folder
    outputs_folder = given_folder
    
    # Check if the outputs folder exists
    if not os.path.exists(outputs_folder):
        return None
    
    # Check if there is a .mp4 file in the outputs folder
    mp4_files = [f for f in os.listdir(outputs_folder) if f.endswith('.mp4')]
    
    # Return the path to the mp4 file if it exists
    if mp4_files:
        return os.path.join(outputs_folder, mp4_files[0])
    else:
        return None


def infer():

    filepath = "./assets/examples/synthetic_1.mp4"
    # Get the current timestamp
    timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")

    output_folder_name = f"results_{timestamp}"

    # Example: Run the inference script (replace with your actual command)
    run_command(f"{sys.executable} inference_keep.py -i={filepath} -o={output_folder_name} --has_aligned --save_video -s=1")

    # Call the function and print the result
    this_infer_folder = os.path.splitext(os.path.basename(filepath))[0]
    joined_path = os.path.join(output_folder_name, this_infer_folder)
    mp4_file_path = check_for_mp4_in_outputs(joined_path)
    print(mp4_file_path)

    print(f"RESULT: {mp4_file_path}")
    
    return result

with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown("# KEEP")
        with gr.Row():
            with gr.Column():
                submit_btn = gr.Button("Submit")
            with gr.Column():
                result = gr.Video()

    submit_btn.click(
        fn = infer, 
        inputs = None,
        outputs = [result],
        show_api=False
    )

demo.queue().launch(show_error=True, show_api=False)