import gradio as gr
# from dotenv import load_dotenv
import os
from huggingface_hub import hf_hub_download
import pandas as pd
import sqlite3

# load_dotenv()

DB_DATASET_ID = os.getenv("DB_DATASET_ID")
DB_NAME = os.getenv("DB_NAME")

cache_path = hf_hub_download(repo_id=DB_DATASET_ID, repo_type='dataset', filename=DB_NAME, token=os.getenv("HF_TOKEN"))

# Model name mappings and metadata

closed_source = [
    'ElevenLabs',
    'Play.HT 2.0',
    'Play.HT 3.0 Mini',
    'PlayDialog',
    'Papla P1',
    'Hume Octave'
]

# Model name mapping, can include models that users cannot vote on
model_names = {
    'styletts2': 'StyleTTS 2',
    'tacotron': 'Tacotron',
    'tacotronph': 'Tacotron Phoneme',
    'tacotrondca': 'Tacotron DCA',
    'speedyspeech': 'Speedy Speech',
    'overflow': 'Overflow TTS',
    'anonymoussparkle': 'Anonymous Sparkle',
    'vits': 'VITS',
    'vitsneon': 'VITS Neon',
    'neuralhmm': 'Neural HMM',
    'glow': 'Glow TTS',
    'fastpitch': 'FastPitch',
    'jenny': 'Jenny',
    'tortoise': 'Tortoise TTS',
    'xtts2': 'Coqui XTTSv2',
    'xtts': 'Coqui XTTS',
    'openvoice': 'MyShell OpenVoice',
    'elevenlabs': 'ElevenLabs',
    'openai': 'OpenAI',
    'hierspeech': 'HierSpeech++',
    'pheme': 'PolyAI Pheme',
    'speecht5': 'SpeechT5',
    'metavoice': 'MetaVoice-1B',
}
model_links = {
    'ElevenLabs': 'https://elevenlabs.io/',
    'Play.HT 2.0': 'https://play.ht/',
    'Play.HT 3.0 Mini': 'https://play.ht/',
    'XTTSv2': 'https://huggingface.co/coqui/XTTS-v2',
    'MeloTTS': 'https://github.com/myshell-ai/MeloTTS',
    'StyleTTS 2': 'https://github.com/yl4579/StyleTTS2',
    'Parler TTS Large': 'https://github.com/huggingface/parler-tts',
    'Parler TTS': 'https://github.com/huggingface/parler-tts',
    'Fish Speech v1.5': 'https://github.com/fishaudio/fish-speech',
    'Fish Speech v1.4': 'https://github.com/fishaudio/fish-speech',
    'GPT-SoVITS': 'https://github.com/RVC-Boss/GPT-SoVITS',
    'WhisperSpeech': 'https://github.com/WhisperSpeech/WhisperSpeech',
    'VoiceCraft 2.0': 'https://github.com/jasonppy/VoiceCraft',
    'PlayDialog': 'https://play.ht/',
    'Kokoro v0.19': 'https://huggingface.co/hexgrad/Kokoro-82M',
    'Kokoro v1.0': 'https://huggingface.co/hexgrad/Kokoro-82M',
    'CosyVoice 2.0': 'https://github.com/FunAudioLLM/CosyVoice',
    'MetaVoice': 'https://github.com/metavoiceio/metavoice-src',
    'OpenVoice': 'https://github.com/myshell-ai/OpenVoice',
    'OpenVoice V2': 'https://github.com/myshell-ai/OpenVoice',
    'Pheme': 'https://github.com/PolyAI-LDN/pheme',
    'Vokan TTS': 'https://huggingface.co/ShoukanLabs/Vokan',
    'Papla P1': 'https://papla.media',
    'Hume Octave': 'https://www.hume.ai'
}


def get_db():
    conn = sqlite3.connect(cache_path)
    return conn

def get_leaderboard(reveal_prelim=False, hide_battle_votes=False, sort_by_elo=True, hide_proprietary=False):
    conn = get_db()
    cursor = conn.cursor()
    
    if hide_battle_votes:
        sql = '''
        SELECT m.name, 
               SUM(CASE WHEN v.username NOT LIKE '%_battle' AND v.vote = 1 THEN 1 ELSE 0 END) as upvote, 
               SUM(CASE WHEN v.username NOT LIKE '%_battle' AND v.vote = -1 THEN 1 ELSE 0 END) as downvote
        FROM model m
        LEFT JOIN vote v ON m.name = v.model
        GROUP BY m.name
        '''
    else:
        sql = '''
        SELECT name, 
               SUM(CASE WHEN vote = 1 THEN 1 ELSE 0 END) as upvote, 
               SUM(CASE WHEN vote = -1 THEN 1 ELSE 0 END) as downvote
        FROM model
        LEFT JOIN vote ON model.name = vote.model
        GROUP BY name
        '''
    
    cursor.execute(sql)
    data = cursor.fetchall()
    df = pd.DataFrame(data, columns=['name', 'upvote', 'downvote'])
    df['name'] = df['name'].replace(model_names).replace('Anonymous Sparkle', 'Fish Speech v1.5')
    
    # Calculate total votes and win rate
    df['votes'] = df['upvote'] + df['downvote']
    df['win_rate'] = (df['upvote'] / df['votes'] * 100).round(1)
    
    # Remove models with no votes
    df = df[df['votes'] > 0]

    # Filter out rows with insufficient votes if not revealing preliminary results
    if not reveal_prelim:
        df = df[df['votes'] > 500]

    ## Calculate ELO SCORE (kept as secondary metric)
    df['elo'] = 1200
    for i in range(len(df)):
        for j in range(len(df)):
            if i != j:
                try:
                    expected_a = 1 / (1 + 10 ** ((df['elo'].iloc[j] - df['elo'].iloc[i]) / 400))
                    expected_b = 1 / (1 + 10 ** ((df['elo'].iloc[i] - df['elo'].iloc[j]) / 400))
                    actual_a = df['upvote'].iloc[i] / df['votes'].iloc[i] if df['votes'].iloc[i] > 0 else 0.5
                    actual_b = df['upvote'].iloc[j] / df['votes'].iloc[j] if df['votes'].iloc[j] > 0 else 0.5
                    df.iloc[i, df.columns.get_loc('elo')] += 32 * (actual_a - expected_a)
                    df.iloc[j, df.columns.get_loc('elo')] += 32 * (actual_b - expected_b)
                except Exception as e:
                    print(f"Error in ELO calculation for rows {i} and {j}: {str(e)}")
                    continue
    df['elo'] = round(df['elo'])

    # Sort based on user preference
    sort_column = 'elo' if sort_by_elo else 'win_rate'
    df = df.sort_values(by=sort_column, ascending=False)
    df['order'] = ['#' + str(i + 1) for i in range(len(df))]
    
    # Select and order columns for display
    df = df[['order', 'name', 'win_rate', 'votes', 'elo']]
    
    # Remove proprietary models if filter is enabled
    if hide_proprietary:
        df = df[~df['name'].isin(closed_source)]
    
    # Convert DataFrame to markdown table with CSS styling
    markdown_table = """
<style>
/* Reset any Gradio table styles */
.leaderboard-table, 
.leaderboard-table th, 
.leaderboard-table td {
    border: none !important;
    border-collapse: separate !important;
    border-spacing: 0 !important;
}

.leaderboard-container {
    background: var(--background-fill-primary);
    border: 1px solid var(--border-color-primary);
    border-radius: 12px;
    padding: 4px;
    margin: 10px 0;
    width: 100%;
    overflow-x: auto;  /* Enable horizontal scroll */
}

.leaderboard-scroll {
    max-height: 600px;
    overflow-y: auto;
    border-radius: 8px;
}

.leaderboard-table {
    width: 100%;
    border-spacing: 0;
    border-collapse: separate;
    font-size: 15px;
    line-height: 1.5;
    table-layout: auto;  /* Allow flexible column widths */
}

.leaderboard-table th {
    background: var(--background-fill-secondary);
    color: var(--body-text-color);
    font-weight: 600;
    text-align: left;
    padding: 12px 16px;
    position: sticky;
    top: 0;
    z-index: 1;
}

.leaderboard-table th:after {
    content: '';
    position: absolute;
    left: 0;
    bottom: 0;
    width: 100%;
    border-bottom: 1px solid var(--border-color-primary);
}

.leaderboard-table td {
    padding: 12px 16px;
    color: var(--body-text-color);
}

.leaderboard-table tr td {
    border-bottom: 1px solid var(--border-color-primary);
}

.leaderboard-table tr:last-child td {
    border-bottom: none;
}

.leaderboard-table tr:hover td {
    background: var(--background-fill-secondary);
}

/* Column-specific styles */
.leaderboard-table .col-rank {
    width: 70px;
    min-width: 70px;  /* Prevent rank from shrinking */
}

.leaderboard-table .col-model {
    min-width: 200px;  /* Minimum width before scrolling */
}

.leaderboard-table .col-winrate {
    width: 100px;
    min-width: 100px;  /* Prevent win rate from shrinking */
}

.leaderboard-table .col-votes {
    width: 100px;
    min-width: 100px;  /* Prevent votes from shrinking */
}

.leaderboard-table .col-arena {
    width: 100px;
    min-width: 100px;  /* Prevent arena score from shrinking */
}

.win-rate {
    display: inline-block;
    font-weight: 600;
    padding: 4px 8px;
    border-radius: 6px;
    min-width: 65px;
    text-align: center;
}

.win-rate-excellent {
    background-color: var(--color-accent);
    color: var(--color-accent-foreground);
}

.win-rate-good {
    background-color: var(--color-accent-soft);
    color: var(--body-text-color);
}

.win-rate-average {
    background-color: var(--background-fill-secondary);
    color: var(--body-text-color);
    border: 1px solid var(--border-color-primary);
}

.win-rate-below {
    background-color: var(--error-background-fill);
    color: var(--body-text-color);
}

.model-link {
    color: var(--body-text-color) !important;
    text-decoration: none !important;
    border-bottom: 2px dashed rgba(128, 128, 128, 0.3);
}

.model-link:hover {
    color: var(--color-accent) !important;
    border-bottom-color: var(--color-accent) !important;
}

.proprietary-badge {
    display: inline-block;
    font-size: 12px;
    padding: 2px 6px;
    border-radius: 4px;
    background-color: var(--background-fill-secondary);
    color: var(--body-text-color);
    margin-left: 6px;
    border: 1px solid var(--border-color-primary);
}

/* New Arena V2 Pointer */
.arena-v2-pointer {
    display: block;
    margin: 20px auto;
    padding: 20px;
    text-align: center;
    border-radius: 12px;
    font-size: 20px;
    font-weight: bold;
    cursor: pointer;
    transition: all 0.3s ease;
    position: relative;
    overflow: hidden;
    text-decoration: none !important;
    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.1);
    max-width: 800px;
    background: linear-gradient(135deg, #FF7B00, #FF5500);
    color: white !important;
    border: none;
}

/* Dark mode adjustments */
@media (prefers-color-scheme: dark) {
    .arena-v2-pointer {
        box-shadow: 0 4px 20px rgba(255, 123, 0, 0.3);
    }
}

.arena-v2-pointer:hover {
    transform: translateY(-5px);
    box-shadow: 0 7px 25px rgba(255, 123, 0, 0.4);
    filter: brightness(1.05);
    color: white !important;
    text-decoration: none !important;
}

.arena-v2-pointer::after {
    content: "→";
    font-size: 24px;
    margin-left: 10px;
    display: inline-block;
    transition: transform 0.3s ease;
}

.arena-v2-pointer:hover::after {
    transform: translateX(5px);
}
</style>

<a href="https://huggingface.co/spaces/TTS-AGI/TTS-Arena-V2" class="arena-v2-pointer" target="_blank">
    Visit the new TTS Arena V2 to vote on the latest models!
</a>

<div class="leaderboard-container">
<div class="leaderboard-scroll">
<table class="leaderboard-table">
<thead>
<tr>
<th class="col-rank">Rank</th>
<th class="col-model">Model</th>
<th class="col-winrate">Win Rate</th>
<th class="col-votes">Votes</th>
""" + ("""<th class="col-arena">Arena Score</th>""" if sort_by_elo else "") + """
</tr>
</thead>
<tbody>
"""
    
    def get_win_rate_class(win_rate):
        if win_rate >= 60:
            return "win-rate-excellent"
        elif win_rate >= 55:
            return "win-rate-good"
        elif win_rate >= 45:
            return "win-rate-average"
        else:
            return "win-rate-below"
    
    for _, row in df.iterrows():
        win_rate_class = get_win_rate_class(row['win_rate'])
        win_rate_html = f'<span class="win-rate {win_rate_class}">{row["win_rate"]}%</span>'
        
        # Add link to model name if available and proprietary badge if closed source
        model_name = row['name']
        original_model_name = model_name
        if model_name in model_links:
            model_name = f'<a href="{model_links[model_name]}" target="_blank" class="model-link">{model_name}</a>'
        
        if original_model_name in closed_source:
            model_name += '<span class="proprietary-badge">Proprietary</span>'
        
        markdown_table += f'''<tr>
            <td class="col-rank">{row['order']}</td>
            <td class="col-model">{model_name}</td>
            <td class="col-winrate">{win_rate_html}</td>
            <td class="col-votes">{row['votes']:,}</td>''' + (
            f'''<td class="col-arena">{int(row['elo'])}</td>''' if sort_by_elo else ""
        ) + "</tr>\n"
    
    markdown_table += "</tbody></table></div></div>"
    return markdown_table

ABOUT = """
# TTS Arena (Legacy)

This is the legacy read-only leaderboard for TTS Arena V1. No new votes are being accepted.

**Please visit the new [TTS Arena](https://huggingface.co/spaces/TTS-AGI/TTS-Arena-V2) to vote!**
"""

CITATION_TEXT = """@misc{tts-arena,
	title        = {Text to Speech Arena},
	author       = {mrfakename and Srivastav, Vaibhav and Fourrier, Clémentine and Pouget, Lucain and Lacombe, Yoach and main and Gandhi, Sanchit},
	year         = 2024,
	publisher    = {Hugging Face},
	howpublished = "\\url{https://huggingface.co/spaces/TTS-AGI/TTS-Arena}"
}"""
FOOTER = f"""
If you reference the Arena in your work, please cite it as follows:

```bibtex
{CITATION_TEXT}
```
"""

with gr.Blocks() as demo:
    gr.Markdown(ABOUT)
    
    with gr.Row():
        with gr.Column():
            reveal_prelim = gr.Checkbox(label="Show preliminary results (< 500 votes)", value=False)
            hide_battle_votes = gr.Checkbox(label="Exclude battle votes", value=False)
        with gr.Column():
            sort_by_elo = gr.Checkbox(label="Sort by Arena Score instead of Win Rate", value=True)
            hide_proprietary = gr.Checkbox(label="Hide proprietary models", value=False)
    
    leaderboard_html = gr.HTML(get_leaderboard())
    
    # Update leaderboard when filters change
    for control in [reveal_prelim, hide_battle_votes, sort_by_elo, hide_proprietary]:
        control.change(
            fn=get_leaderboard,
            inputs=[reveal_prelim, hide_battle_votes, sort_by_elo, hide_proprietary],
            outputs=leaderboard_html
        )
    
    gr.Markdown(FOOTER)

demo.launch()