import av
import os
import torch
import tempfile
import shutil
import atexit
import subprocess
import gradio as gr

from convert import convert_video


def get_video_length_av(video_path):
    with av.open(video_path) as container:
        stream = container.streams.video[0]
        if container.duration is not None:
            duration_in_seconds = float(container.duration) / av.time_base
        else:
            duration_in_seconds = stream.duration * stream.time_base

    return duration_in_seconds


def get_video_dimensions(video_path):
    with av.open(video_path) as container:
        video_stream = container.streams.video[0]
        width = video_stream.width
        height = video_stream.height

    return width, height


def get_free_memory_gb():
    gpu_index = torch.cuda.current_device()
    gpu_properties = torch.cuda.get_device_properties(gpu_index)

    total_memory = gpu_properties.total_memory
    allocated_memory = torch.cuda.memory_allocated(gpu_index)

    free_memory = total_memory - allocated_memory
    return free_memory / 1024**3


def cleanup_temp_directories():
    print("Deleting temporary files")
    for temp_dir in temp_directories:
        try:
            shutil.rmtree(temp_dir)
        except FileNotFoundError:
            print(f"Could not delete directory {temp_dir}")


def ffmpeg_remux_audio(source_video_path, dest_video_path, output_path):
    # Build the ffmpeg command to extract audio and remux into another video
    command = [
        "ffmpeg",
        "-i",
        dest_video_path,  # Input destination video file
        "-i",
        source_video_path,  # Input source video file (for the audio)
        "-c:v",
        "copy",  # Copy the video stream as is
        "-c:a",
        "copy",  # Copy the audio stream as is
        "-map",
        "0:v:0",  # Map the video stream from the destination file
        "-map",
        "1:a:0",  # Map the audio stream from the source file
        output_path,  # Specify the output file path
    ]

    try:
        # Run the ffmpeg command
        subprocess.run(command, check=True)
    except subprocess.CalledProcessError as e:
        # Handle errors during the subprocess execution
        print(f"An error occurred: {e}")
        return dest_video_path

    return output_path


def inference(video):
    if get_video_length_av(video) > 30:
        raise gr.Error("Length of video cannot be over 30 seconds")
    if get_video_dimensions(video) > (1920, 1920):
        raise gr.Error("Video resolution must not be higher than 1920x1080")

    temp_dir = tempfile.mkdtemp()
    temp_directories.append(temp_dir)

    output_composition = temp_dir + "/matted_video.mp4"
    convert_video(
        model,  # The loaded model, can be on any device (cpu or cuda).
        input_source=video,  # A video file or an image sequence directory.
        downsample_ratio=0.25,  # [Optional] If None, make downsampled max size be 512px.
        output_composition=output_composition,  # File path if video; directory path if png sequence.
        output_alpha=None,  # [Optional] Output the raw alpha prediction.
        output_foreground=None,  # [Optional] Output the raw foreground prediction.
        output_video_mbps=4,  # Output video mbps. Not needed for png sequence.
        seq_chunk=12,  # Process n frames at once for better parallelism.
        num_workers=1,  # Only for image sequence input. Reader threads.
        progress=True,  # Print conversion progress.
    )

    resulting_video = f"{temp_dir}/matted_{os.path.split(video)[1]}"

    return ffmpeg_remux_audio(video, output_composition, resulting_video)


if __name__ == "__main__":
    temp_directories = []
    atexit.register(cleanup_temp_directories)

    model = torch.hub.load(
        "PeterL1n/RobustVideoMatting", "mobilenetv3", trust_repo=True
    )

    if torch.cuda.is_available():
        free_memory = get_free_memory_gb()
        concurrency_count = int(free_memory // 7)
        print(f"Using GPU with concurrency: {concurrency_count}")
        print(f"Available video memory: {free_memory} GB")
        model = model.cuda()
    else:
        print("Using CPU")
        concurrency_count = 1

    with gr.Blocks(title="Robust Video Matting") as block:
        gr.Markdown("# Robust Video Matting")
        gr.Markdown(
            "Gradio demo for Robust Video Matting. To use it, simply upload your video, or click one of the examples to load them. Read more at the links below."
        )
        with gr.Row():
            inp = gr.Video(label="Input Video", sources=["upload"], include_audio=True)
            out = gr.Video(label="Output Video")
        btn = gr.Button("Run")
        btn.click(inference, inputs=inp, outputs=out)

        gr.Examples(
            examples=[["example.mp4"]],
            inputs=[inp],
        )
        gr.HTML(
            "<p style='text-align: center'><a href='https://arxiv.org/abs/2108.11515'>Robust High-Resolution Video Matting with Temporal Guidance</a> | <a href='https://github.com/PeterL1n/RobustVideoMatting'>Github Repo</a></p>"
        )

    block.queue(api_open=False, max_size=5, concurrency_count=concurrency_count).launch(
        share=False
    )