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 = f"acap_{Path(video_path).name}"
    new_video_path = f"{VIDEO_DIRECTORY}/{new_video}"
    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"{VIDEO_DIRECTORY}/{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 Exception(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 Exception(f"Failed to download video, error code: {error_code}")
    
    return output_path

def wrap_html(yt_url):
    return f"""<h1><center>Original Video: <a href='{yt_url}'>{yt_url}</a></center><h1>"""

# 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.rsplit("=", 1)[-1]

    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(yt_url)


def load_mystery_video(gr_request: gr.Request):

    video_paths = list(Path(VIDEO_DIRECTORY).glob("*.mp4"))

    selected_video = video_paths[random.randrange(len(video_paths))]

    file_name = selected_video.name.replace("acap_", "").split(".", 1)[0]
    selected_audio = Path(VIDEO_DIRECTORY) / f"{file_name}.m4a"
    yt_url = f"https://www.youtube.com/watch?v={file_name}"

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

    return selected_video, selected_audio, wrap_html(yt_url)


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

with gr.Blocks(theme=theme) as demo:

    with gr.Row():
        header = gr.Markdown(markdown_text)
    with gr.Row().style():
        
        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()