import os
import gradio as gr
import json
import numpy as np
from datetime import datetime
import os
import sys
import librosa

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


# 模拟歌曲生成函数
def generate_song(description, lyric, prompt_audio=None):
    # 这里模拟生成过程 - 实际应用中替换为你的模型调用
    print(f"Generating song with description: {description}")
    print(f"Lyrics provided: {lyric}")
    if prompt_audio is not None:
        print("Using prompt audio for generation")
    
    # 从文件中加载示例音频
    audio_path = "./sample/example.mp3"
    audio_data, sample_rate = librosa.load(audio_path, sr=None)  # 保持原始采样率
    
    
    # 创建输入配置的JSON
    input_config = {
        "description": description,
        "lyric": lyric,
        "has_prompt_audio": prompt_audio is not None,
        "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"):
                    text_prompt = gr.Textbox(
                    label="Text Prompt",
                    placeholder="Enter the Text Prompt, eg: emotional piano pop",
                    )
            
            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],
        outputs=[output_audio, output_json]
    )
    

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