import datetime

import gradio as gr
import json
import requests
import time


def create_task(prompt: str, style: str, title: str, api_key: str):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    url = "https://api.kie.ai/api/v1/generate"
    payload = {
        "prompt": prompt,
        "style": style,
        "title": title,
        "customMode": False,
        "instrumental": True,
        "model": "V3_5",
        "callBackUrl": "https://your-app.com/callback"
    }
    try:
        response = requests.post(url, headers=headers, data=json.dumps(payload), timeout=60)
        result = response.json()

        if response.ok and result.get('code') == 200:
            print(f"任务ID: {result['data']['taskId']}")
            return True, result['data']['taskId']
        else:
            print(f"请求失败: {result.get('msg', '未知错误')}")
            return False, result.get('msg', "Internal Server Error")
    except Exception:
        pass
    return False, "Internal Server Error"


def check_task_status(task_id, api_key):
    url = f"https://api.kie.ai/api/v1/generate/record-info?taskId={task_id}"
    headers = {"Authorization": f"Bearer {api_key}"}
    for _ in range(3):
        try:
            response = requests.get(url, headers=headers)
            result = response.json()

            if response.ok and result.get('code') == 200:
                task_data = result['data']
                status = task_data['status']

                response_data = task_data['response']

                if status == 'SUCCESS':
                    return "SUCCESS", response_data["sunoData"]
                elif status == 'FIRST_SUCCESS':
                    return "FIRST_SUCCESS", response_data["sunoData"]
                elif status == 'TEXT_SUCCESS':
                    return "TEXT_SUCCESS", task_data
                elif status == 'PENDING':
                    return "PENDING", None
                else:
                    return "FAILED", response_data.get('errorMessage', 'Internal Server Error')
            else:
                return "FAILED", result.get('msg', 'Internal Server Error')
        except Exception:
            pass
        time.sleep(5)
    return "FAILED", 'Internal Server Error'


def format_duration(seconds):
    """格式化时长显示"""
    minutes = int(seconds // 60)
    seconds = int(seconds % 60)
    return f"{minutes:02d}:{seconds:02d}"


def create_music_card(music_data, index):
    """创建音乐卡片的HTML"""
    if music_data is None:
        return f"""
        <div class="music-card empty-card">
            <div class="card-content">
                <div class="empty-icon">🎵</div>
                <div class="empty-text">Waiting for music #{index} to generate...</div>
            </div>
        </div>
        """

    # 格式化时长
    duration = format_duration(music_data.get('duration', 0))

    # 格式化创建时间
    create_time = music_data.get('createTime', '')
    if create_time:
        # 如果是毫秒级时间戳，先转成秒
        if create_time > 1e12:
            create_time = create_time / 1000

        # 转换为 datetime 对象
        dt = datetime.datetime.fromtimestamp(create_time)
        # 格式化输出
        formatted_time = dt.strftime("%Y-%m-%d %H:%M:%S")
    else:
        formatted_time = ""

    print(formatted_time)
    return f"""
    <div class="music-card">
        <div class="card-header">
            <img src="{music_data.get('imageUrl', '')}" alt="Album cover" class="album-cover">
            <div class="track-number">#{index}</div>
        </div>
        <div class="card-content">
            <h3 class="music-title">{music_data.get('title', 'Untitled')}</h3>
            <div class="music-tags">
                <span class="tag-icon">🏷️</span>
                <span class="tag-text">{music_data.get('tags', 'No tags')}</span>
            </div>
            <div class="music-duration">
                <span class="duration-icon">⏱️</span>
                <span class="duration-text">{duration}</span>
            </div>
            <audio controls class="audio-player">
                <source src="{music_data.get('audioUrl', '')}" type="audio/mpeg">
                Your browser does not support audio playback
            </audio>
            <div class="create-time">Generated: {formatted_time}</div>
        </div>
    </div>
    """


def process_music_generation(prompt, style, title, api_key, progress=gr.Progress()):
    yield (
        create_music_card(None, 1),
        create_music_card(None, 2),
        f"🚀 Creating music generation task..."
    )
    """处理音乐生成的主函数"""
    if not api_key or not api_key.strip():
        yield (gr.update(),  # 保持音乐卡片不变
               gr.update(),  # 保持音乐卡片不变
               f"❌ Please enter API Key")
        return "❌ Please enter API Key"

    if not prompt or not prompt.strip():
        yield (gr.update(),  # 保持音乐卡片不变
               gr.update(),  # 保持音乐卡片不变
               f"❌ Please enter music description")
        return "❌ Please enter music description"

    try:
        # 步骤1：创建任务
        success, task_id = create_task(prompt.strip(), style.strip(), title.strip(), api_key.strip())

        if not success:
            yield (gr.update(),  # 保持音乐卡片不变
                   gr.update(),  # 保持音乐卡片不变
                   f"❌ Failed to create task: {task_id}")
            return f"❌ Failed to create task"

        yield (gr.update(),  # 保持音乐卡片不变
               gr.update(),  # 保持音乐卡片不变
               f"📋 Starting music generation")

        # 步骤2：轮询检查状态
        start_time = time.time()
        first_music_shown = False
        music_data_1 = None
        music_data_2 = None
        last_api_check = 0
        api_check_interval = 8  # API每8秒检查一次

        while time.time() - start_time < 600:  # 最多等待10分钟
            elapsed = int(time.time() - start_time)
            progress_value = min(0.2 + (elapsed / 600) * 0.7, 0.9)

            # 每8秒检查一次API状态，避免频繁请求
            if time.time() - last_api_check >= api_check_interval:
                status, data = check_task_status(task_id, api_key.strip())
                last_api_check = time.time()

            if status == 'FIRST_SUCCESS' and not first_music_shown:
                # 第一首音乐生成完成
                progress(progress_value, desc="🎵 First music generated!")
                if data and len(data) > 0:
                    music_data_1 = data[0]
                    first_music_shown = True
                    # 立即更新第一首音乐的显示
                    yield (
                        create_music_card(music_data_1, 1),
                        create_music_card(None, 2),
                        f"✅ First music generated, generating second..."
                    )
                    continue

            elif status == 'SUCCESS':
                # 两首音乐都生成完成
                progress(0.95, desc="🎉 All music generated!")
                if data and len(data) > 0:
                    music_data_1 = data[0] if len(data) > 0 else music_data_1
                    music_data_2 = data[1] if len(data) > 1 else None

                    yield (
                        create_music_card(music_data_1, 1),
                        create_music_card(music_data_2, 2) if music_data_2 else create_music_card(None, 2),
                        f"✅ Successfully generated {len(data)} music tracks!"
                    )
                    break

            elif status == 'FAILED':
                yield (
                    gr.update(),  # 保持音乐卡片不变
                    gr.update(),  # 保持音乐卡片不变
                    f"❌ Generation failed: {data}"
                )
                break

            if not first_music_shown:
                yield (
                    gr.update(),  # 保持音乐卡片不变
                    gr.update(),  # 保持音乐卡片不变
                    f"⏳ Generating music... (waited {elapsed}s)"
                )
            elif first_music_shown:
                # 第一首音乐已生成，只更新状态文本
                yield (
                    gr.update(),  # 保持音乐卡片不变
                    gr.update(),  # 保持音乐卡片不变
                    f"✅ First music generated, generating second... (waited {elapsed}s)"
                )

            time.sleep(1)

        # 超时
        if time.time() - start_time >= 600:
            yield (
                gr.update(),  # 保持音乐卡片不变
                gr.update(),  # 保持音乐卡片不变
                "⚠️ Generation timeout, please try again later"
            )

    except Exception as e:
        yield (
            gr.update(),  # 保持音乐卡片不变
            gr.update(),  # 保持音乐卡片不变
            f"❌ Error occurred: {str(e)}"
        )
        return "❌ Error occurred: {str(e)}"


# CSS样式
css = """
.gradio-container {
    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
    min-height: 100vh;
}
.header-container {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    padding: 2.5rem;
    border-radius: 24px;
    margin-bottom: 2.5rem;
    box-shadow: 0 20px 60px rgba(102, 126, 234, 0.25);
}
.logo-text {
    font-size: 3.5rem;
    font-weight: 900;
    color: white;
    text-align: center;
    margin: 0;
    letter-spacing: -2px;
    text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.3);
}
.subtitle {
    color: rgba(255, 255, 255, 0.9);
    text-align: center;
    font-size: 1.2rem;
    margin-top: 0.5rem;
}
.main-content {
    background: rgba(255, 255, 255, 0.95);
    backdrop-filter: blur(20px);
    border-radius: 24px;
    padding: 2.5rem;
    box-shadow: 0 10px 40px rgba(0, 0, 0, 0.08);
}
.mode-indicator {
    background: rgba(255, 255, 255, 0.3);
    backdrop-filter: blur(10px);
    border-radius: 12px;
    padding: 0.5rem 1rem;
    margin-top: 1rem;
    text-align: center;
    font-weight: 600;
    color: white;
}

.info-box {
    background: linear-gradient(135deg, #e3f2fd 0%, #bbdefb 100%);
    border-radius: 12px;
    padding: 1.5rem;
    margin-bottom: 1.5rem;
    border-left: 4px solid #2196f3;
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1);
}

.info-box strong {
    color: #1976d2;
    font-size: 1.1rem;
}

.gr-button-primary {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    color: white !important;
    font-weight: 700 !important;
    font-size: 1.1rem !important;
    padding: 1.2rem 2rem !important;
    border-radius: 14px !important;
    text-transform: uppercase;
    letter-spacing: 1px;
    width: 100%;
    margin-top: 1rem !important;
    transition: all 0.3s ease !important;
}

.gr-button-primary:hover {
    background: linear-gradient(135deg, #5a67d8 0%, #6b46c1 100%) !important;
    transform: translateY(-2px);
    box-shadow: 0 6px 20px rgba(0, 0, 0, 0.15) !important;
}

/* 音乐卡片样式 */
.music-cards-container {
    display: flex;
    gap: 1.5rem;
    margin-top: 1rem;
    flex-wrap: wrap;
}

.music-card {
    background: white;
    border-radius: 16px;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.1);
    overflow: hidden;
    transition: all 0.3s ease;
    flex: 1;
    min-width: 300px;
    animation: slideIn 0.5s ease-out;
}

@keyframes slideIn {
    from {
        opacity: 0;
        transform: translateY(20px);
    }
    to {
        opacity: 1;
        transform: translateY(0);
    }
}

.music-card:hover {
    transform: translateY(-5px);
    box-shadow: 0 8px 30px rgba(0, 0, 0, 0.15);
}

.empty-card {
    border: 2px dashed #e0e0e0;
    background: #fafafa;
    min-height: 400px;
    display: flex;
    align-items: center;
    justify-content: center;
}

.empty-icon {
    font-size: 3rem;
    margin-bottom: 1rem;
    opacity: 0.3;
}

.empty-text {
    color: #9e9e9e;
    font-size: 1rem;
}

.card-header {
    position: relative;
    height: 200px;
    overflow: hidden;
}

.album-cover {
    width: 100%;
    height: 100%;
    object-fit: cover;
}

.track-number {
    position: absolute;
    top: 10px;
    right: 10px;
    background: rgba(0, 0, 0, 0.7);
    color: white;
    padding: 5px 12px;
    border-radius: 20px;
    font-weight: bold;
    font-size: 0.9rem;
}

.card-content {
    padding: 1.5rem;
}

.music-title {
    font-size: 1.4rem;
    font-weight: 700;
    color: #2c3e50;
    margin: 0 0 1rem 0;
}

.music-tags {
    display: flex;
    align-items: center;
    margin-bottom: 0.8rem;
    color: #7f8c8d;
}

.tag-icon, .duration-icon {
    margin-right: 0.5rem;
    font-size: 1rem;
}

.tag-text {
    font-size: 0.9rem;
    font-style: italic;
}

.music-duration {
    display: flex;
    align-items: center;
    margin-bottom: 1rem;
    color: #7f8c8d;
}

.duration-text {
    font-weight: 600;
}

.audio-player {
    width: 100%;
    margin: 1rem 0;
    border-radius: 8px;
}

.create-time {
    font-size: 0.8rem;
    color: #95a5a6;
    text-align: right;
    margin-top: 1rem;
}

/* 输入框样式 */
.gr-input, .gr-textarea {
    background: #ffffff !important;
    border: 2px solid #e1e8ed !important;
    border-radius: 10px !important;
    color: #374151 !important;
    font-size: 1rem !important;
    padding: 0.75rem 1rem !important;
    transition: all 0.3s ease !important;
}

.gr-input:focus, .gr-textarea:focus {
    border-color: #667eea !important;
    box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
}

/* 下拉框样式 */
.gr-dropdown {
    background: #ffffff !important;
    border: 2px solid #e1e8ed !important;
    border-radius: 10px !important;
}

label {
    color: #4a5568 !important;
    font-weight: 600 !important;
    font-size: 0.9rem !important;
    margin-bottom: 0.5rem !important;
}

/* 示例样式 */
.examples-holder {
    margin-top: 2rem;
    padding: 1.5rem;
    background: #f8f9fa;
    border-radius: 12px;
}

/* 隐藏页脚 */
footer {
    display: none !important;
}
"""

# 创建Gradio界面
with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
    with gr.Column(elem_classes="header-container"):
        gr.HTML("""
            <h1 class="logo-text">🎵 Suno AI Music API Free Online Test</h1>
            <p class="subtitle">Supported by the latest Suno v4.5 and Suno v4.5+ models</p>
            <div class="mode-indicator">
                💡 Test Suno AI Music API for free on Hugging Face Spaces. Generate music, edit vocals, and create tracks with the latest Suno v4.5 and Suno v4.5+ models. 
            </div>
        """)

    with gr.Column(elem_classes="main-content"):
        # 信息提示框
        gr.HTML("""
            <div class="info-box">
                <strong>Supported Features of Suno API:</strong><br>
                • <b>Generate Music:</b> Create original songs with or without lyrics.<br>
                • <b>Extend Music:</b> Seamlessly extend your tracks to any length.<br>
                • <b>Generate Lyrics:</b> Create unique lyrics from text prompts.<br>
                • <b>Add Vocals/Instrumentals:</b> Add vocals or instrumental accompaniments to your tracks.<br>
                • <b>Separate Vocals:</b> Isolate vocals or instruments for editing.<br>
                • <b>Convert to WAV:</b> Export music in high-quality WAV format.<br>
                <strong>Usage Instructions: </strong><br>
                • Get your Suno API Key <a href="https://kie.ai/playground/suno" target="_blank" style="color: #1976d2;">👉 here 👈</a><br>
                • Select your model: Choose from Suno v4.5+, Suno v4.5, Suno v4, or Suno v3.5.<br>
                • Choose Custom Mode (optional), specify if you want to generate instrumental music only, and enter your music style description or lyrics.<br>
                • Click Generate and wait ~1–2 minutes for processing.<br>
            </div>
        """)

        with gr.Row():
            # 左侧输入区域
            with gr.Column(scale=1):
                # 第一块：API Key
                with gr.Group():
                    gr.Markdown("### 🔑 API Configuration")
                    api_key = gr.Textbox(
                        label="API Key",
                        placeholder="Please enter your API Key",
                        type="password"
                    )

                # 第二块：音乐信息
                with gr.Group():
                    gr.Markdown("### 🎵 Music Information")
                    title = gr.Textbox(
                        label="Music Title",
                        placeholder="Give your music a name",
                        value=""
                    )

                    style = gr.Textbox(
                        label="Music Style",
                        placeholder="Enter your desired music style",
                        interactive=True
                    )

                    prompt = gr.Textbox(
                        label="Music Description",
                        placeholder="Describe the music you want, e.g.: An energetic electronic dance track with futuristic synthesizer sounds...",
                        lines=3,
                    )

                generate_btn = gr.Button(
                    "🎼 Generate Music",
                    variant="primary",
                    size="lg"
                )

            # 右侧输出区域
            with gr.Column(scale=2):
                gr.Markdown("### 🎵 Generation Results")

                # 两个音乐展示区域
                with gr.Row(elem_classes="music-cards-container"):
                    music_display_1 = gr.HTML(
                        value=create_music_card(None, 1),
                        elem_id="music-1"
                    )
                    music_display_2 = gr.HTML(
                        value=create_music_card(None, 2),
                        elem_id="music-2"
                    )

                # 状态显示
                status_text = gr.Textbox(
                    label="Generation Status",
                    interactive=False,
                    value="Ready! Please enter music description and click generate button...",
                    lines=2
                )

        # 音乐创意示例 - 放在页面最下面
        gr.Examples(
            examples=[
                ["Summer Beach", "Pop", "A light and cheerful summer beach music with fresh guitar and upbeat rhythm"],
                ["Burning Heart", "Rock", "Passionate rock music with powerful electric guitar and drums"],
                ["Neon Night", "Electronic", "Futuristic electronic music with cyberpunk style"],
                ["Afternoon Time", "Jazz", "Gentle jazz piano piece, perfect for café ambiance"],
                ["Heroic Epic", "Classical", "Epic orchestral music, grand and majestic, perfect for film soundtrack"],
            ],
            inputs=[title, style, prompt],
            label="🎵 Music Creative Examples"
        )

    # 绑定生成事件
    generate_btn.click(
        fn=process_music_generation,
        inputs=[prompt, style, title, api_key],
        outputs=[music_display_1, music_display_2, status_text]
    )

# 启动应用
if __name__ == "__main__":
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860  # 使用不同的端口
    )