import os
import shutil
import subprocess
import gradio as gr

# execute a CLI command
def execute_command(command: str) -> None:
    subprocess.run(command, check=True)


def infer(video_frames, masks_frames, project_name):


    # Create the directory if it doesn't exist
    my_video_directory = f"{project_name}"
    if not os.path.exists(my_video_directory):
        os.makedirs(my_video_directory)

    # Assuming 'images' is a list of image file paths
    for idx, image in enumerate(video_frames):
        # Get the base file name (without path) from the original location
        image_name = os.path.basename(image.name)
    
        # Construct the destination path in the working directory
        destination_path = os.path.join(my_video_directory, image_name)
    
        # Copy the image from the original location to the working directory
        shutil.copy(image.name, destination_path)
    
        # Print the image name and its corresponding save path
        print(f"Image {idx + 1}: {image_name} copied to {destination_path}")


    # Create the directory if it doesn't exist
    my_masks_directory = f"{project_name}_masks"
    if not os.path.exists(my_masks_directory):
        os.makedirs(my_masks_directory)

    # Assuming 'images' is a list of image file paths
    for idx, image in enumerate(masks_frames):
        # Get the base file name (without path) from the original location
        image_name = os.path.basename(image.name)
    
        # Construct the destination path in the working directory
        destination_path = os.path.join(my_masks_directory, image_name)
    
        # Copy the image from the original location to the working directory
        shutil.copy(image.name, destination_path)
    
        # Print the image name and its corresponding save path
        print(f"Image {idx + 1}: {image_name} copied to {destination_path}")

    #video_frames_folder = "inputs/object_removal/bmx-trees"
    #masks_folder = "inputs/object_removal/bmx-trees_mask"

    video_frames_folder = f"{my_video_directory}"
    masks_folder = f"{my_masks_directory}"
    
    # Create the "results" folder if it doesn't exist
    output_folder = "results"
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)


    #bmx_trees_folder = os.path.join(output_folder, "bmx-trees")
    
    command = [
      f"python", 
      f"inference_propainter.py",
      f"--video={video_frames_folder}",
      f"--mask={masks_folder}",
      f"--output={output_folder}"
    ]

    execute_command(command)

    # Get the list of files in the "results" folder
    result_files = os.listdir(output_folder)

    # Print the content of the "results" folder
    print(f"Contents of the {output_folder} folder:")
    for item in result_files:
        print(item)

    # List the content of the "bmx-trees" folder within "results"
    results_folder = os.path.join(output_folder, f"{project_name}")
    results_folder_content = [os.path.join(results_folder, item) for item in os.listdir(results_folder)]

    print(f"Contents of the {results_folder} folder:")
    for item in results_folder_content:
        print(item)

    return "done", results_folder_content[0], results_folder_content[1]

css="""
#col-container{
    margin: 0 auto;
    max-width: 840px;
    text-align: left;
}
"""
    
with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML("""
        <h2 style="text-align: center;">ProPainter</h2>
        <p style="text-align: center;">
            [ICCV 2023] ProPainter: Improving Propagation and Transformer for Video Inpainting <br />
            <a href="https://github.com/sczhou/ProPainter" target="_blank">code</a> | <a href="https://shangchenzhou.com/projects/ProPainter/" target="_blank">project page</a>
        </p>
                """)

        with gr.Row():
            with gr.Column():
                project_name = gr.Textbox(label="Name your project", value="my-new-project")
                video_frames = gr.File(label="Video frames", file_types=["image"], file_count="multiple")
                masks_frames = gr.File(label="Masks frames", file_types=["image"], file_count="multiple")
        
                submit_btn = gr.Button("Submit")

            with gr.Column():
                result = gr.Textbox(label="Result")
                res_masked = gr.Video()
                res_files = gr.Video()

        
            
    submit_btn.click(fn=infer, inputs=[video_frames, masks_frames, project_name], outputs=[result, res_masked, res_files])
    
demo.queue(max_size=12).launch()