import os
from scipy.io.wavfile import write
import gradio as gr
import yt_dlp

os.makedirs("ytdl", exist_ok=True)

def separator(audio):
    os.makedirs("output", exist_ok=True)
    write('test.wav', audio[0], audio[1])
    result = os.system("audio-separator test.wav --model_filename model_bs_roformer_ep_317_sdr_12.9755.ckpt --output_dir=./output --output_format=wav")

    files = [
        "./output/test_(Instrumental)_model_bs_roformer_ep_317_sdr_12.wav",
        "./output/test_(Vocals)_model_bs_roformer_ep_317_sdr_12.wav"
    ]

    for file in files:
        if not os.path.isfile(file):
            print(f"File not found: {file}")
        else:
            print(f"File exists: {file}")

    return files

def download_audio(url):
    ydl_opts = {
        'format': 'bestaudio/best',
        'outtmpl': 'ytdl/%(title)s.%(ext)s',
        'postprocessors': [{
            'key': 'FFmpegExtractAudio',
            'preferredcodec': 'mp3',
            'preferredquality': '192',
        }],
    }

    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info_dict = ydl.extract_info(url, download=True)
        file_path = ydl.prepare_filename(info_dict).rsplit('.', 1)[0] + '.mp3'
        return file_path

def download_and_show(url):
    file_path = download_audio(url)
    return file_path

with gr.Blocks() as demo:
    with gr.Tab("Separator"):
        audio_input = gr.Audio(type="numpy", label="Input")
        output1 = gr.Audio(type="filepath", label="Stem 1", interactive=False)
        output2 = gr.Audio(type="filepath", label="Stem 2", interactive=False)
        submit_button = gr.Button("Separate")
        submit_button.click(
            separator, 
            inputs=[audio_input],
            outputs=[output1, output2]
        )
    with gr.Tab("Download YouTube audio for separation"):
        url = gr.Textbox(label="URL to YouTube")
        download = gr.Button("Download")
        download.click(
            download_and_show,
            inputs=[url],
            outputs=[audio_input]
        )

demo.launch()