import os
import librosa
import soundfile
import gradio as gr
from scipy.io import wavfile
import subprocess
from pathlib import Path
import time
import random
from datetime import datetime
from scipy.io.wavfile import write
from pydub import AudioSegment
high_quality = True

'''
os.system("git clone https://github.com/google/REAPER.git")
os.system("cd REAPER")
os.system("mkdir build")
os.system("cd build")
os.system("cmake ..")
os.system("make")
os.system("cd ..")
os.system("cd ..")
'''

def svc(audio_src, audio_ref, singing_check):
  wav1, sr1 = librosa.load(audio_src, sr=24000)
  soundfile.write("audio_source.wav", wav1, sr1)
  wav2, sr2 = librosa.load(audio_ref, sr=24000)
  soundfile.write("audio_reference.wav", wav2, sr2)
  if singing_check == True:
    os.system("python infer.py --src_wav_path audio_source.wav --ref_wav_path audio_reference.wav --out_path output.wav")

  else:
    os.system("python infer.py --src_wav_path audio_source.wav --ref_wav_path audio_reference.wav --out_path output.wav --speech_enroll")

  return "output.wav"

def youtube_downloader(
    video_identifier,
    start_time,
    end_time,
    is_full_song,
    output_filename="track.wav",
    num_attempts=5,
    url_base="",
    quiet=False,
    force=True,
):
    if is_full_song:
        ydl_opts = {
            'noplaylist': True,
            'format': 'bestaudio/best',
            'postprocessors': [{
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'wav',
            }],
            "outtmpl": 'dl_audio/youtube_audio',
        }
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([video_identifier])
        audio_path = "dl_audio/youtube_audio.wav"
        return audio_path

    else:
        output_path = Path(output_filename)
        if output_path.exists():
            if not force:
                return output_path
            else:
                output_path.unlink()

        quiet = "--quiet --no-warnings" if quiet else ""
        command = f"""
            yt-dlp {quiet} -x --audio-format wav -f bestaudio -o "{output_filename}" --download-sections "*{start_time}-{end_time}" "{url_base}{video_identifier}"  # noqa: E501
        """.strip()

        attempts = 0
        while True:
            try:
                _ = subprocess.check_output(command, shell=True, stderr=subprocess.STDOUT)
            except subprocess.CalledProcessError:
                attempts += 1
                if attempts == num_attempts:
                    return None
            else:
                break

        if output_path.exists():
            return output_path
        else:
            return None

def audio_separated(audio_input, progress=gr.Progress()):
    # start progress
    progress(progress=0, desc="Starting...")
    time.sleep(0.1)

    # check file input
    if audio_input is None:
        # show progress
        for i in progress.tqdm(range(100), desc="Please wait..."):
            time.sleep(0.01)

        return (None, None, 'Please input audio.')

    # create filename
    filename = str(random.randint(10000,99999))+datetime.now().strftime("%d%m%Y%H%M%S")

    # progress
    progress(progress=0.10, desc="Please wait...")

    # make dir output
    os.makedirs("output", exist_ok=True)

    # progress
    progress(progress=0.20, desc="Please wait...")

    # write
    if high_quality:
        write(filename+".wav", audio_input[0], audio_input[1])
    else:
        write(filename+".mp3", audio_input[0], audio_input[1])

    # progress
    progress(progress=0.50, desc="Please wait...")

    # demucs process
    if high_quality:
        command_demucs = "python3 -m demucs --two-stems=vocals -d cpu "+filename+".wav -o output"
    else:
        command_demucs = "python3 -m demucs --two-stems=vocals --mp3 --mp3-bitrate 128 -d cpu "+filename+".mp3 -o output"

    os.system(command_demucs)

    # progress
    progress(progress=0.70, desc="Please wait...")

    # remove file audio
    if high_quality:
        command_delete = "rm -v ./"+filename+".wav"
    else:
        command_delete = "rm -v ./"+filename+".mp3"

    os.system(command_delete)

    # progress
    progress(progress=0.80, desc="Please wait...")

    # progress
    for i in progress.tqdm(range(80,100), desc="Please wait..."):
        time.sleep(0.1)

    if high_quality:
        return "./output/htdemucs/"+filename+"/vocals.wav","./output/htdemucs/"+filename+"/no_vocals.wav","Successfully..."
    else:
        return "./output/htdemucs/"+filename+"/vocals.mp3","./output/htdemucs/"+filename+"/no_vocals.mp3","Successfully..."

# mix vocal and non-vocal
def mix(audio1, audio2):
  sound1 = AudioSegment.from_file(audio1)
  sound2 = AudioSegment.from_file(audio2)
  length = len(sound1)
  mixed = sound1[:length].overlay(sound2)

  mixed.export("song.wav", format="wav")

  return "song.wav"

app = gr.Blocks()

with app:
  gr.Markdown("# <center>🥳💕🎶 NeuCo AI歌手，无需训练、一键翻唱</center>")
  gr.Markdown("## <center>🌟 只需10秒音频，一键翻唱任意歌手的任意歌曲！Powered by [NeuCoSVC](https://github.com/thuhcsi/NeuCoSVC)</center>")
  gr.Markdown("### <center>🌊 更多精彩应用，敬请关注[滔滔AI](http://www.talktalkai.com)；滔滔AI，为爱滔滔！💕</center>")
  with gr.Row():
    with gr.Column():
      ydl_url_input  = gr.Textbox(label="音乐视频网址(可直接填写相应的BV号)", value = "https://www.bilibili.com/video/BV...")
      with gr.Row():
        start = gr.Number(value=0, label="起始时间 (秒)")
        end = gr.Number(value=15, label="结束时间 (秒)")
        check_full = gr.Checkbox(label="是否上传整首歌曲", info="若勾选则不需要填写起止时间", value=False)
      ydl_url_submit = gr.Button("提取声音文件吧", variant="primary")
      as_audio_submit = gr.Button("去除背景音吧", variant="primary")

    with gr.Column():
      ydl_audio_output = gr.Audio(label="歌曲原声", interactive=True)
      as_audio_input = ydl_audio_output
      as_audio_vocals = gr.Audio(label="歌曲人声部分", type="filepath", interactive=False)
      as_audio_no_vocals = gr.Audio(label="歌曲伴奏部分", type="filepath", visible=False)
      as_audio_message = gr.Textbox(label="Message", visible=False)

  ydl_url_submit.click(fn=youtube_downloader, inputs=[ydl_url_input, start, end, check_full], outputs=[ydl_audio_output])
  as_audio_submit.click(fn=audio_separated, inputs=[as_audio_input], outputs=[as_audio_vocals, as_audio_no_vocals, as_audio_message], show_progress=True, queue=True)

  with gr.Row():
    with gr.Column():
      inp_src = as_audio_vocals
      inp_ref = gr.Audio(label="请上传您喜欢的参考音频", type="filepath")
      check_singing = gr.Checkbox(label="参考音频是否为歌曲", info="勾选：参考音频是歌曲；不勾选：参考音频是语音", value=True)
      btn1 = gr.Button("一键歌声转换吧！", variant="primary")
      btn2 = gr.Button("加入歌曲伴奏吧！", variant="primary")

    with gr.Column():
      out_svc = gr.Audio(label="为您生成的专属原声", type="filepath", interactive=False)
      new_song = gr.Audio(label="AI歌手+伴奏", type="filepath")

  btn1.click(svc, [inp_src, inp_ref, check_singing], out_svc)
  btn2.click(fn=mix, inputs=[out_svc, as_audio_no_vocals], outputs=[new_song])

  gr.Markdown("### <center>注意❗：请不要生成会对个人以及组织造成侵害的内容，此程序仅供科研、学习及个人娱乐使用。</center>")
  gr.HTML('''
      <div class="footer">
                  <p>🌊🏞️🎶 - 江水东流急，滔滔无尽声。 明·顾璘
                  </p>
      </div>
  ''')

app.queue().launch(show_error=True)