import gradio as gr
import pandas as pd
import re
from pathlib import Path

# Column order matching the original
COL_ORDER = [
    "Model Name",
    "WorldScore-Static",
    "WorldScore-Dynamic",
    "Camera Control",
    "Object Control",
    "Content Alignment",
    "3D Consistency",
    "Photometric Consistency",
    "Style Consistency",
    "Subjective Quality",
    "Motion Accuracy",
    "Motion Magnitude",
    "Motion Smoothness",
    "Model Type",
    "Ability",
    "Sampled by",
    "Evaluated by",
    "Accessibility",
    "Date",
]

# Numeric columns for highlighting
NUMERIC_COLS = {
    "WorldScore-Static",
    "WorldScore-Dynamic",
    "Camera Control",
    "Object Control",
    "Content Alignment",
    "3D Consistency",
    "Photometric Consistency",
    "Style Consistency",
    "Subjective Quality",
    "Motion Accuracy",
    "Motion Magnitude",
    "Motion Smoothness",
}


def parse_markdown_link(text):
    """Parse markdown link format [text](url) and return HTML link"""
    match = re.match(r'^\[(.*?)\]\((.*?)\)$', str(text))
    if match:
        text_content = match.group(1)
        url = match.group(2)
        # Return HTML link for Gradio to render
        return f'<a href="{url}" target="_blank" style="color: #1d4ed8; text-decoration: none; font-weight: 600;">{text_content}</a>'
    return str(text)


def load_and_process_data():
    """Load CSV and process data"""
    csv_path = Path("leaderboard.csv")
    if not csv_path.exists():
        return pd.DataFrame()
    
    df = pd.read_csv(csv_path)
    
    # Reorder columns according to COL_ORDER (only include columns that exist)
    existing_cols = [col for col in COL_ORDER if col in df.columns]
    df = df[existing_cols]
    
    # Parse markdown links in Model Name column and convert to HTML
    if "Model Name" in df.columns:
        df["Model Name"] = df["Model Name"].apply(parse_markdown_link)
    
    # Convert numeric columns to float for proper sorting
    for col in NUMERIC_COLS:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors='coerce')
    
    # Sort by WorldScore-Static descending by default
    if "WorldScore-Static" in df.columns:
        df = df.sort_values("WorldScore-Static", ascending=False, na_position='last')
    
    # Format numeric columns to 2 decimal places
    for col in NUMERIC_COLS:
        if col in df.columns:
            df[col] = df[col].apply(lambda x: f"{x:.2f}" if pd.notna(x) else "")
    
    return df


def create_leaderboard():
    """Create the Gradio interface"""
    df = load_and_process_data()
    
    # Custom CSS to match the original styling
    css = """
    .gradio-container {
        background-color: #0b1120;
        color: #e5e7eb;
        font-family: system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
    }
    .intro-section {
        padding: 24px 32px;
        color: #e5e7eb;
    }
    .intro-section h1 {
        font-size: 2rem;
        font-weight: 800;
        margin-bottom: 8px;
        color: #ffffff;
    }
    .intro-links {
        margin-bottom: 12px;
        font-size: 1rem;
    }
    .intro-links a {
        color: #60a5fa;
        text-decoration: none;
        font-weight: 600;
        margin-right: 8px;
    }
    .intro-links a:hover {
        text-decoration: underline;
    }
    .intro-text {
        color: #d1d5db;
        font-size: 1rem;
        margin-top: 6px;
        line-height: 1.5rem;
    }
    .divider {
        margin: 20px 0;
        border: none;
        height: 1px;
        background-color: #374151;
    }
    """
    
    # Create intro HTML
    intro_html = """
    <div class="intro-section">
        <h1>WorldScore Leaderboard</h1>
        <div class="intro-links">
            <a href="https://arxiv.org/abs/2504.00983" target="_blank">Paper</a> |
            <a href="https://haoyi-duan.github.io/WorldScore/" target="_blank">Website</a> |
            <a href="https://github.com/haoyi-duan/WorldScore" target="_blank">Code</a> |
            <a href="https://huggingface.co/datasets/Howieeeee/WorldScore" target="_blank">Dataset</a> |
            <a href="https://github.com/haoyi-duan/WorldScore?tab=readme-ov-file#leaderboard" target="_blank">Join Leaderboard</a>
        </div>
        <p class="intro-text">
            🏆 Welcome to the leaderboard of <b>WorldScore</b>, the first unified evaluation benchmark for world generation.
        </p>
        <hr class="divider">
    </div>
    """
    
    with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="blue")) as demo:
        gr.HTML(intro_html)
        
        # Create dataframe component with sorting enabled
        if len(df.columns) > 0:
            # Set datatype for Model Name column to markdown to render HTML links
            datatypes = ["markdown" if col == "Model Name" else "str" for col in df.columns]
            dataframe = gr.Dataframe(
                value=df,
                label="",
                interactive=False,
                wrap=True,
                column_widths=["auto"] * len(df.columns),
                height=600,
                datatype=datatypes,
            )
        else:
            gr.Markdown("No data available. Please ensure leaderboard.csv exists.")
        
    return demo


if __name__ == "__main__":
    demo = create_leaderboard()
    demo.launch()