import os
import torch
import librosa
import binascii
import warnings
import midi2audio  # MIDI 파일을 WAV 파일로 변환
import numpy as np
import pytube as pt  # YouTube 비디오를 오디오로 다운로드
import gradio as gr
import soundfile as sf
from transformers import Pop2PianoForConditionalGeneration, Pop2PianoProcessor

# 디렉토리 생성
yt_video_dir = "./yt_dir"  # 유튜브 비디오 다운로드 경로
outputs_dir = "./midi_wav_outputs"  # 출력 파일 경로
os.makedirs(outputs_dir, exist_ok=True)
os.makedirs(yt_video_dir, exist_ok=True)

# 모델 설정
device = "cuda" if torch.cuda.is_available() else "cpu"
model = Pop2PianoForConditionalGeneration.from_pretrained("sweetcocoa/pop2piano").to(device)
processor = Pop2PianoProcessor.from_pretrained("sweetcocoa/pop2piano")
composers = model.generation_config.composer_to_feature_token.keys()

# 유튜브 비디오에서 오디오 추출 함수
def get_audio_from_yt_video(yt_link):
    try:
        yt = pt.YouTube(yt_link)
        t = yt.streams.filter(only_audio=True)
        filename = os.path.join(yt_video_dir, binascii.hexlify(os.urandom(8)).decode() + ".mp4")
        t[0].download(filename=filename)
    except:
        warnings.warn(f"Video Not Found at {yt_link}")
        filename = None
    
    return filename, filename

# 모델 추론 함수
def inference(file_uploaded, composer):
    waveform, sr = librosa.load(file_uploaded, sr=None) 
    inputs = processor(audio=waveform, sampling_rate=sr, return_tensors="pt").to(device)
    model_output = model.generate(input_features=inputs["input_features"], composer=composer)
    tokenizer_output = processor.batch_decode(token_ids=model_output.to("cpu"), feature_extractor_output=inputs.to("cpu"))["pretty_midi_objects"]

    return prepare_output_file(tokenizer_output, sr)    

# 출력 파일 준비 함수
def prepare_output_file(tokenizer_output, sr):
    output_file_name = "output_" + binascii.hexlify(os.urandom(8)).decode()
    midi_output = os.path.join(outputs_dir, output_file_name + ".mid")
    tokenizer_output[0].write(midi_output)
    wav_output = midi_output.replace(".mid", ".wav")
    midi2audio.FluidSynth().midi_to_audio(midi_output, wav_output)
    
    return wav_output, wav_output, midi_output

# Gradio UI 설정
block = gr.Blocks(theme="Taithrah/Minimal")

with block:
    gr.HTML(
        """
        <div style="text-align: center; max-width: 800px; margin: 0 auto;">
            <h1 style="font-weight: 900; margin-bottom: 12px;">
                🎹 Pop2Piano : 피아노 커버곡 생성기 🎹
            </h1>
            <p style="margin-bottom: 12px; font-size: 90%">
                Pop2Piano 데모: 팝 오디오 기반 피아노 커버곡 생성. <br>
                작곡가(편곡자)를 선택하고 팝 오디오를 업로드하거나 유튜브 링크를 입력한 후 생성 버튼을 클릭하세요.
            </p>
        </div>
        """
    )
    with gr.Group():
        with gr.Row():
            with gr.Column():
                file_uploaded = gr.Audio(label="오디오 업로드", type="filepath")
            with gr.Column():
                with gr.Row():
                    yt_link = gr.Textbox(label="유튜브 링크를 입력하세요.", autofocus=True, lines=3)
                    yt_btn = gr.Button("유튜브 링크에서 오디오를 다운 받습니다.", size="lg")
                yt_audio_path = gr.Audio(label="유튜브 동영상에서 추출한 오디오", interactive=False)
                yt_btn.click(get_audio_from_yt_video, inputs=[yt_link], outputs=[yt_audio_path, file_uploaded])

    with gr.Group():
        with gr.Column():
            composer = gr.Dropdown(label="편곡자", choices=composers, value="composer1")
            generate_btn = gr.Button("나만의 피아노 커버곡 만들기🎹🎵")
        with gr.Row():
            wav_output2 = gr.File(label="나만의 피아노 커버곡을 다운로드 (.wav)")
            wav_output1 = gr.Audio(label="나만의 피아노 커버곡 듣기")
            midi_output = gr.File(label="생성한 midi 파일 다운로드 (.mid)")
            generate_btn.click(
                inference, 
                inputs=[file_uploaded, composer], 
                outputs=[wav_output1, wav_output2, midi_output])

block.launch(debug=False)