from gradio_client import Client
import yt_dlp  # pip install yt-dlp

import subprocess

import os
from pathlib import Path
import gradio as gr
import shutil
import random
from theme import theme

VIDEO_DIRECTORY = "videos"


client = Client("abidlabs/music-separation")


def acapellify(audio_path):
    result = client.predict(audio_path, api_name="/predict")
    return result[0]


# based on https://github.com/gradio-app/gradio/blob/bebfb72b353a4280155cf7070441fc476ac10172/guides/06_client-libraries/fastapi-app-with-the-gradio-client.md
def process_video(video_path):
    old_audio = os.path.basename(video_path).split(".")[0] + ".m4a"
    subprocess.run(
        ["ffmpeg", "-y", "-i", video_path, "-vn", "-acodec", "copy", old_audio]
    )

    new_audio = acapellify(old_audio)

    new_video_path = f"{VIDEO_DIRECTORY}/{Path(video_path).name}"
    subprocess.call(
        [
            "ffmpeg",
            "-y",
            "-i",
            video_path,
            "-i",
            new_audio,
            "-map",
            "0:v",
            "-map",
            "1:a",
            "-c:v",
            "copy",
            "-c:a",
            "aac",
            "-strict",
            "experimental",
            new_video_path,
        ]
    )

    # remove old audio and video
    os.remove(old_audio)
    os.remove(video_path)

    new_audio_path = f"{VIDEO_DIRECTORY}/{old_audio}"
    shutil.move(new_audio, new_audio_path)

    return new_video_path, new_audio_path


# filename default value will return name of video on youtube
def download_yt_url(url: str, filename: str = "%(title)s", format="mp4"):
    output_path = f"{filename}.{format}"

    # restrict video length so one user doesn't take up all our bandwidth
    def video_filter(info):
        MAX_DURATION = 10 * 60
        duration = info.get("duration")
        if duration and duration > MAX_DURATION:
            raise gr.Error(
                f"The video is too long at {duration}s, choose a video less than {MAX_DURATION}s"
            )

    ydl_opts = {
        "match_filter": video_filter,
        "format": f"{format}/bestaudio",
        "outtmpl": output_path,
    }

    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        error_code = ydl.download([url])
        # info = ydl.extract_info(url, extra_info={"output_path": output_path})

    if error_code:
        raise gr.Error(f"Failed to download video, error code: {error_code}")

    return output_path


def wrap_html(filename):
    yt_url = f"https://www.youtube.com/watch?v={filename}"

    return f"""<h1><center>Original Video</center></h1> 
        <center><a href='{yt_url}'>{yt_url}</a></center>
        <br>
        <center><iframe width="100%" height="315" src="https://www.youtube.com/embed/{filename}" title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share" allowfullscreen></iframe></center>
    """


# ideally yt_url should be validated
def acapellify_url(yt_url, gr_request: gr.Request):
    # example filename: https://www.youtube.com/watch?v=TasKo5HHWb4 -> TasKo5HHWb4
    filename = yt_url.split("=", 1)[-1].split("=", 1)[0]

    video_path = download_yt_url(yt_url, filename)

    new_video_path, new_audio_path = process_video(video_path)
    return new_video_path, new_audio_path, wrap_html(filename)


def load_mystery_video(gr_request: gr.Request):
    video_paths = list(Path(VIDEO_DIRECTORY).glob("*.mp4"))
    n_videos = len(video_paths)

    if not n_videos:
        raise gr.Error(
            "No videos archived yet. Enter a Youtube URL to add the first video!"
        )

    selected_video = video_paths[random.randrange(n_videos)]

    filename = selected_video.name.split(".", 1)[0]
    selected_audio = Path(VIDEO_DIRECTORY) / f"{filename}.m4a"

    selected_video = str(selected_video)
    selected_audio = str(selected_audio)

    return selected_video, selected_audio, wrap_html(filename)


with open("header.md", "r") as markdown_file:
    markdown_text = markdown_file.read()

with gr.Blocks(theme=theme, css="footer {visibility: hidden}") as demo:
    with gr.Row():
        header = gr.Markdown(markdown_text)
    with gr.Row().style(equal_height=True):
        with gr.Column(scale=0.4, variant="panel"):
            input_url = gr.Textbox(label="Youtube URL")

            process_video_btn = gr.Button("Acapellify", variant="primary")
            mystery_btn = gr.Button("Mysterious Video")

            static_path_display = gr.HTML(label="Output File Paths", visible=True)

        with gr.Column(scale=0.6, variant="panel"):
            output_video = gr.Video(label="Acapellified Video")
            output_audio = gr.Audio(label="Acapellified Audio")

    process_video_btn.click(
        fn=acapellify_url,
        inputs=[input_url],
        outputs=[output_video, output_audio, static_path_display],
    )
    mystery_btn.click(
        fn=load_mystery_video,
        inputs=[],
        outputs=[output_video, output_audio, static_path_display],
    )

if __name__ == "__main__":
    os.makedirs(VIDEO_DIRECTORY, exist_ok=True)
    demo.launch()