import os
import gradio as gr
import json
import numpy as np
from datetime import datetime
import os
import yaml
import sys
import librosa
import time
import os.path as op
APP_DIR = op.dirname(op.abspath(__file__))

from download import download_model
# 下载模型
download_model(APP_DIR)
print("Successful downloaded model.")

from levo_inference import LeVoInference

MODEL = LeVoInference(op.join(APP_DIR, "conf/infer.yaml"))

EXAMPLE_DESC = """female, dark, pop, sad, piano and drums, the bpm is 125."""
EXAMPLE_LYRICS = """
[intro-short]

[verse]
夜晚的街灯闪烁.
我漫步在熟悉的角落.
回忆像潮水般涌来.
你的笑容如此清晰.
在心头无法抹去.
那些曾经的甜蜜.
如今只剩我独自回忆.

[bridge]
手机屏幕亮起.
是你发来的消息.
简单的几个字.
却让我泪流满面.
曾经的拥抱温暖.
如今却变得遥远.
我多想回到从前.
重新拥有你的陪伴.

[chorus]
回忆的温度还在.
你却已不在.
我的心被爱填满.
却又被思念刺痛.
R&B的节奏奏响.
我的心却在流浪.
没有你的日子.
我该如何继续向前.

[outro-short]
""".strip()

with open('conf/vocab.yaml', 'r', encoding='utf-8') as file:
    STRUCTS = yaml.safe_load(file)


# 模拟歌曲生成函数
def generate_song(description, lyric, prompt_audio=None, cfg_coef=None, temperature=None, top_k=None, progress=gr.Progress(track_tqdm=True)):
    global MODEL
    global STRUCTS
    params = {'cfg_coef':cfg_coef, 'temperature':temperature, 'top_k':top_k}
    params = {k:v for k,v in params.items() if v is not None}
    sample_rate = MODEL.cfg.sample_rate

    # 生成过程
    print(f"Generating song with description: {description}")
    print(f"Lyrics provided: {lyric}")
    
    # 适配lyric格式
    lyric = lyric.replace("\n\n", " ; ")
    for s in STRUCTS:
        lyric = lyric.replace(f"{s}\n", f"{s} ")
    lyric = lyric.replace("\n", "")
    lyric = lyric.replace(". ; ", " ; ")
    
    # 适配prompt 
    if prompt_audio is not None:
        print("Using prompt audio for generation")
    else:
        prompt_audio = op.join(APP_DIR, 'sample/prompt.wav')

    progress(0.0, "Start Generation")
    start = time.time()
    
    audio_data = MODEL(lyric, description, prompt_audio, params).cpu().permute(1, 0).float().numpy()

    end = time.time()
    
    # 创建输入配置的JSON
    input_config = {
        "description": description,
        "lyric": lyric,
        "prompt_audio": prompt_audio,
        "params": params,
        "inference_duration": end - start,
        "timestamp": datetime.now().isoformat(),
    }
    
    return (sample_rate, audio_data), json.dumps(input_config, indent=2)

# 创建Gradio界面
with gr.Blocks(title="LeVo Demo Space") as demo:
    gr.Markdown("# 🎵 LeVo Demo Space")
    gr.Markdown("Demo interface for the LeVo song generation model. Provide a description, lyrics, and optionally an audio prompt to generate a custom song.")
    
    with gr.Row():
        with gr.Column():
            description = gr.Textbox(
                label="Song Description",
                placeholder="Describe the style, mood, and characteristics of the song...",
                lines=1,
                max_lines=2,
                value=EXAMPLE_DESC,
            )
            lyric = gr.Textbox(
                label="Lyrics",
                placeholder="Enter the lyrics for the song...",
                lines=5,
                max_lines=8,
                value=EXAMPLE_LYRICS,
            )

            with gr.Tabs(elem_id="extra-tabs"):
                with gr.Tab("Audio Prompt"):
                    prompt_audio = gr.Audio(
                        label="Prompt Audio (Optional)",
                        type="filepath",
                        elem_id="audio-prompt"
                    )
                with gr.Tab("Advanced Config"):
                    cfg_coef = gr.Slider(
                        label="CFG Coefficient",
                        minimum=0.1,
                        maximum=3.0,
                        step=0.1,
                        value=1.5,
                        interactive=True,
                        elem_id="cfg-coef",
                    )
                    temperature = gr.Slider(
                        label="Temperature",
                        minimum=0.1,
                        maximum=2.0,
                        step=0.1,
                        value=1.0,
                        interactive=True,
                        elem_id="temperature",
                    )
                    top_k = gr.Slider(
                        label="Top-K",
                        minimum=1,
                        maximum=100,
                        step=1,
                        value=50,
                        interactive=True,
                        elem_id="top_k",
                    )
            generate_btn = gr.Button("Generate Song", variant="primary")
        
        with gr.Column():
            output_audio = gr.Audio(label="Generated Song", type="numpy")
            output_json = gr.JSON(label="Input Configuration")
    
    # 示例按钮
    examples = gr.Examples(
        examples=[
            ["An uplifting pop song with catchy melodies"],
            ["Melancholic piano ballad"],
        ],
        inputs=[description],
        label="Description examples"
    )

    examples = gr.Examples(
        examples=[
            ["Shine bright like the stars above\nYou're the one that I'm dreaming of"],
            ["The rain keeps falling on my window pane\nReminding me of love that's gone away"],
        ],
        inputs=[lyric],
        label="Lyrics examples"
    )
    
    # 生成按钮点击事件

    generate_btn.click(
        fn=generate_song,
        inputs=[description, lyric, prompt_audio, cfg_coef, temperature, top_k],
        outputs=[output_audio, output_json]
    )
    

# 启动应用
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)