import gradio as gr
from utils import *
import os
from leaderboard import build_leaderboard

# Custom CSS styling similar to the reference space
css = """
/* Custom CSS that works with Ocean theme */
.song-arena-header {
    text-align: center;
    padding: 20px;
    margin-bottom: 20px;
    border-radius: 10px;
}
.song-arena-logo {
    max-width: 150px;
    border-radius: 10px;
    box-shadow: 0 4px 8px rgba(0,0,0,0.3);
}
.song-arena-title {
    font-size: 28px;
    margin-bottom: 10px;
}
.song-arena-subtitle {
    margin-bottom: 15px;
}
.song-arena-description {
    font-size: 16px;
    margin: 0;
}
/* Resource links styling */
.resource-links {
    display: flex;
    justify-content: center;
    flex-wrap: wrap;
    gap: 8px;
    margin-bottom: 25px;
}
.resource-link {
    background-color: #222222;
    color: #4aedd6;
    border: 1px solid #333333;
    padding: 8px 16px;
    border-radius: 20px;
    margin: 5px;
    text-decoration: none;
    display: inline-block;
    font-weight: 500;
    box-shadow: 0 2px 4px rgba(0, 0, 0, 0.3);
    transition: all 0.2s ease;
}
.resource-link:hover {
    background-color: #333333;
    transform: translateY(-2px);
    box-shadow: 0 3px 6px rgba(0, 0, 0, 0.4);
    transition: all 0.2s ease;
}
.resource-link-icon {
    margin-right: 5px;
}
/* Footer styling */
.song-arena-footer {
    text-align: center;
    margin-top: 30px;
    padding: 15px;
}
/* Research warning */
.research-warning {
    background-color: #ff4c4c;
    color: white;
    text-align: center;
    padding: 10px;
    border-radius: 8px;
    font-weight: bold;
    margin: 15px 0;
}
"""

def build_demo():
    last_updated_path = 'results/last_updated.txt'
    last_updated_str = ''
    if os.path.exists(last_updated_path):
        with open(last_updated_path, 'r') as f:
            date = f.read()
            last_updated_str = f"\nLast Updated: {date}"
    
    with gr.Blocks(css=css) as demo:
        state = gr.State()
        
        # Header and Paper Information
        gr.HTML(
            """
            <div class="song-arena-header">
                <div style="display: flex; justify-content: center; margin-bottom: 20px;">
                    <img src="https://i.postimg.cc/3Jx3yZ5b/real-vs-fake-sonics-w-logo.jpg" class="song-arena-logo">
                </div>
                <h1 class="song-arena-title">🎵 Song Arena</h1>
                <h3 class="song-arena-subtitle">SONICS: Synthetic Or Not - Identifying Counterfeit Songs | ICLR 2025 [Poster]</h3>
                <p class="song-arena-description">
                    Can you detect if a song is AI-generated or real? Test your skills!
                </p>
            </div>
            """
        )
        
        # Resource Links
        gr.HTML(
            """
            <div class="resource-links">
                <a href="https://openreview.net/forum?id=PY7KSh29Z8" target="_blank" class="resource-link">
                    <span class="resource-link-icon">📄</span>Paper
                </a>
                <a href="https://huggingface.co/datasets/awsaf49/sonics" target="_blank" class="resource-link">
                    <span class="resource-link-icon">🎵</span>Dataset
                </a>
                <a href="https://huggingface.co/collections/awsaf49/sonics-spectttra-67bb6517b3920fd18e409013" target="_blank" class="resource-link">
                    <span class="resource-link-icon">🤖</span>Models
                </a>
                <a href="https://arxiv.org/abs/2408.14080" target="_blank" class="resource-link">
                    <span class="resource-link-icon">🔬</span>ArXiv
                </a>
                <a href="https://github.com/awsaf49/sonics" target="_blank" class="resource-link">
                    <span class="resource-link-icon">💻</span>GitHub
                </a>
            </div>
            """
        )
        
        # Research purpose warning
        gr.HTML(
            """
            <div class="research-warning">
                ⚠️ FOR RESEARCH PURPOSES ONLY ⚠️
            </div>
            """
        )
        
        with gr.Tab("Song Generation", id=0):
            gr.Markdown("# 🎵 Song Arena\nCan you detect if a song is AI-generated or real?")
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### 👇 Click on **New Round** for a song!")
                    with gr.Row():
                        with gr.Column():
                            model_selector = gr.Markdown("", visible=True)
                    with gr.Row():
                        with gr.Column():
                            # heard_btn = gr.Button(
                            #     value="🚩 Heard this song before", visible=False, interactive=False
                            # )
                            audio = gr.Audio(show_download_button=False)
                    with gr.Row():
                        real_btn = gr.Button(value="😃 Real", visible=False, interactive=False)
                        fake_btn = gr.Button(
                            value="🤖 Fake", visible=False, interactive=False
                        )
                        btn_list = [fake_btn, real_btn]
                        fake_btn.click(
                            fake_last_response,
                            inputs=[state],
                            outputs=[fake_btn, real_btn, model_selector]
                        )
                        real_btn.click(
                            real_last_response,
                            inputs=[state],
                            outputs=[fake_btn, real_btn, model_selector]
                        )
                    # heard_btn.click(generate_songs, state, [state, audio, model_selector])
                    new_round_button = gr.Button("🎵 New Round 🎵")
                    new_round_button.click(generate_songs, state, [state, audio, model_selector]).then(
                        enable_buttons_side_by_side,
                        inputs=None,
                        outputs=btn_list
                    )
            
            # How It Works Section
            with gr.Accordion("How It Works", open=True):
                gr.Markdown("""
                    ### The SONICS Game
                    
                    This interactive game challenges you to distinguish between real songs created by humans and those generated by AI. Your votes help us understand how well humans can detect AI-generated music.
                    
                    ### Instructions:
                    - Click "New Round" to get a random song
                    - Listen carefully
                    - Vote whether you think it's real (human-created) or fake (AI-generated)
                    - Your scores will be tracked on the leaderboard
                """)
                
        with gr.Tab("Leaderboard", id=1):
            gr.Markdown("# 🏆 Leaderboard 🏆" + last_updated_str)
            gr.Dataframe(
                headers=[
                    "🤖/ 😃",
                    "🗳️ Num Votes",
                    "📊 Sensitivity",
                    "📊 Specificity",
                    "📊 F-1",
                ],
                datatype=[
                    "str",
                    "number",
                    "number",
                ],
                value=build_leaderboard().values,
                max_height=1200,
                column_widths=[200, 100, 100, 100, 100, 100, 100],
                wrap=False,
            )
        
        # Footer
        gr.HTML(
            """
            <div class="song-arena-footer">
                <p>SONICS: Synthetic Or Not - Identifying Counterfeit Songs | ICLR 2025</p>
                <p style="font-size: 12px;">For research purposes only</p>
            </div>
            """
        )
        
    return demo

if __name__ == "__main__":
    demo = build_demo()
    demo.queue(max_size=20).launch(server_name="0.0.0.0")