"""
ACE-Step v1.5 音樂生成器
基於 Gradio 的 Web UI，調用 ACE-Step 官方 API 進行音樂生成
"""

import os
import base64
import tempfile
import httpx
import gradio as gr
from dotenv import load_dotenv

# 載入環境變數
load_dotenv()

# API 配置
API_BASE_URL = os.environ.get("ACEMUSIC_API_URL", "https://api.acemusic.ai")
API_KEY = os.environ.get("ACEMUSIC_API_KEY", "")


def generate_music(
    prompt: str,
    lyrics: str,
    vocal_language: str,
    duration: float,
    bpm: int,
    instrumental: bool,
    sample_mode: bool,
    use_format: bool,
    progress=gr.Progress()
):
    """
    調用 ACE-Step API 生成音樂
    
    Args:
        prompt: 音樂描述
        lyrics: 歌詞（選填）
        vocal_language: 人聲語言
        duration: 音訊時長（秒）
        bpm: 節拍速度
        instrumental: 是否為純音樂
        sample_mode: 是否使用自然語言模式
        use_format: 是否使用 LLM 格式化增強
        progress: Gradio 進度條
    
    Returns:
        tuple: (音訊檔案路徑, 元數據文字, 錯誤訊息)
    """
    if not prompt.strip():
        return None, "", "請輸入音樂描述"
    
    if not API_KEY:
        return None, "", "未設定 API Key，請在環境變數中設定 ACEMUSIC_API_KEY"
    
    progress(0.1, desc="準備請求...")
    
    # 構建訊息內容
    if lyrics.strip():
        # 歌詞模式：使用標籤格式
        content = f"<prompt>{prompt}</prompt>\n<lyrics>{lyrics}</lyrics>"
    elif sample_mode:
        # 自然語言模式
        content = prompt
    else:
        # 標籤模式
        content = f"<prompt>{prompt}</prompt>"
    
    # 構建請求體
    request_body = {
        "messages": [
            {"role": "user", "content": content}
        ],
        "audio_config": {
            "duration": duration,
            "vocal_language": vocal_language,
        },
        "sample_mode": sample_mode,
        "use_format": use_format,
    }
    
    # 添加可選參數
    if bpm > 0:
        request_body["audio_config"]["bpm"] = bpm
    
    if instrumental:
        request_body["audio_config"]["instrumental"] = True
    
    progress(0.2, desc="發送請求到 API...")
    
    try:
        # 發送請求
        with httpx.Client(timeout=600) as client:
            response = client.post(
                f"{API_BASE_URL}/v1/chat/completions",
                headers={
                    "Content-Type": "application/json",
                    "Authorization": f"Bearer {API_KEY}"
                },
                json=request_body
            )
        
        progress(0.8, desc="處理回應...")
        
        # 檢查回應狀態
        if response.status_code == 401:
            return None, "", "API Key 無效或未授權"
        elif response.status_code == 429:
            return None, "", "服務繁忙，請稍後再試"
        elif response.status_code == 503:
            return None, "", "模型尚未初始化，請稍後再試"
        elif response.status_code == 504:
            return None, "", "生成逾時，請嘗試縮短時長"
        elif response.status_code != 200:
            return None, "", f"API 錯誤: {response.status_code} - {response.text}"
        
        # 解析回應
        result = response.json()
        
        # 提取音訊資料
        choices = result.get("choices", [])
        if not choices:
            return None, "", "API 回應格式錯誤：缺少 choices"
        
        message = choices[0].get("message", {})
        audio_list = message.get("audio", [])
        
        if not audio_list:
            return None, "", "API 回應中沒有音訊資料"
        
        # 獲取音訊 URL
        audio_url = audio_list[0].get("audio_url", {}).get("url", "")
        
        if not audio_url:
            return None, "", "音訊 URL 為空"
        
        # 解碼 Base64 音訊
        progress(0.9, desc="解碼音訊...")
        
        # 處理 Data URL 格式: data:audio/mpeg;base64,<base64_data>
        if audio_url.startswith("data:"):
            b64_data = audio_url.split(",", 1)[1]
        else:
            b64_data = audio_url
        
        audio_bytes = base64.b64decode(b64_data)
        
        # 儲存為臨時檔案
        with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as f:
            f.write(audio_bytes)
            audio_path = f.name
        
        # 提取元數據
        content_text = message.get("content", "音樂生成成功")
        
        progress(1.0, desc="完成！")
        
        return audio_path, content_text, ""
        
    except httpx.TimeoutException:
        return None, "", "請求逾時，請稍後再試"
    except httpx.RequestError as e:
        return None, "", f"網路錯誤: {str(e)}"
    except Exception as e:
        return None, "", f"發生錯誤: {str(e)}"


# 建立 Gradio 介面
def create_ui():
    """建立 Gradio UI"""
    
    with gr.Blocks(
        title="ACE-Step 音樂生成器",
        theme=gr.themes.Soft(),
        css="""
        .container { max-width: 1200px; margin: auto; }
        .header { text-align: center; margin-bottom: 20px; }
        .output-audio { margin-top: 20px; }
        """
    ) as demo:
        
        gr.HTML("""
        <div class="header">
            <h1>🎵 ACE-Step v1.5 音樂生成器</h1>
            <p>使用 AI 生成高品質音樂 - 支援多種風格與語言</p>
        </div>
        """)
        
        with gr.Row():
            # 左側：輸入區
            with gr.Column(scale=1):
                gr.Markdown("### 📝 輸入設定")
                
                prompt_input = gr.Textbox(
                    label="音樂描述",
                    placeholder="描述您想要的音樂風格、情緒、場景...\n例如：輕快的華語流行歌曲，鋼琴伴奏",
                    lines=3,
                    max_lines=5
                )
                
                lyrics_input = gr.Textbox(
                    label="歌詞（選填）",
                    placeholder="輸入歌詞，支援結構標記...\n\n[Verse 1]\n第一段歌詞\n\n[Chorus]\n副歌歌詞",
                    lines=5,
                    max_lines=10
                )
                
                with gr.Accordion("進階設定", open=False):
                    sample_mode_check = gr.Checkbox(
                        label="自然語言模式",
                        value=False,
                        info="啟用後，系統會自動生成提示詞和歌詞"
                    )
                    
                    use_format_check = gr.Checkbox(
                        label="LLM 格式化增強",
                        value=False,
                        info="使用 LLM 增強提示詞和歌詞"
                    )
            
            # 右側：參數設定
            with gr.Column(scale=1):
                gr.Markdown("### ⚙️ 參數設定")
                
                language_dropdown = gr.Dropdown(
                    label="人聲語言",
                    choices=[
                        ("中文", "zh"),
                        ("英文", "en"),
                        ("日文", "ja"),
                        ("韓文", "ko"),
                        ("自動", "auto")
                    ],
                    value="zh",
                    interactive=True
                )
                
                duration_slider = gr.Slider(
                    label="時長（秒）",
                    minimum=10,
                    maximum=120,
                    value=30,
                    step=5,
                    interactive=True
                )
                
                bpm_input = gr.Slider(
                    label="BPM（選填，0 表示自動決定）",
                    minimum=0,
                    maximum=200,
                    value=0,
                    step=1,
                    interactive=True
                )
                
                instrumental_check = gr.Checkbox(
                    label="純音樂（無人聲）",
                    value=False,
                    interactive=True
                )
        
        # 生成按鈕
        generate_btn = gr.Button(
            "🎵 生成音樂",
            variant="primary",
            size="lg"
        )
        
        # 輸出區
        gr.Markdown("### 🎧 生成結果")
        
        with gr.Row():
            with gr.Column(scale=2):
                audio_output = gr.Audio(
                    label="生成的音訊",
                    type="filepath",
                    interactive=False
                )
            
            with gr.Column(scale=1):
                metadata_output = gr.Textbox(
                    label="元數據",
                    lines=10,
                    max_lines=15,
                    interactive=False
                )
        
        error_output = gr.Textbox(
            label="錯誤訊息",
            visible=False
        )
        
        # 綁定事件
        generate_btn.click(
            fn=generate_music,
            inputs=[
                prompt_input,
                lyrics_input,
                language_dropdown,
                duration_slider,
                bpm_input,
                instrumental_check,
                sample_mode_check,
                use_format_check
            ],
            outputs=[
                audio_output,
                metadata_output,
                error_output
            ]
        )
        
        # 使用說明
        gr.Markdown("""
        ---
        ### 📖 使用說明
        
        **輸入模式：**
        1. **標籤模式**（推薦）：直接輸入音樂描述，系統會自動識別
        2. **歌詞模式**：輸入歌詞（支援 `[Verse]`、`[Chorus]` 等標記）
        3. **自然語言模式**：用自然語言描述，讓 AI 自動生成
        
        **參數說明：**
        - **時長**：生成音訊的長度（10-120 秒）
        - **BPM**：節拍速度，留空由系統自動決定
        - **純音樂**：勾選後不會生成人聲
        
        ---
        Powered by [ACE-Step](https://acemusic.ai) | API: `api.acemusic.ai`
        """)
    
    return demo


# 主程式進入點
if __name__ == "__main__":
    demo = create_ui()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )