import requests
from bs4 import BeautifulSoup
import pandas as pd
import gradio as gr
import io

BASE_URL = "https://scale.com/leaderboard"

LEADERBOARDS = {
    "Coding": "/coding",
    "Adversarial Robustness": "/adversarial_robustness",
    "Instruction Following": "/instruction_following",
    "Math": "/math"
}

def scrape_leaderboard(leaderboard):
    url = BASE_URL + LEADERBOARDS[leaderboard]
    
    response = requests.get(url)
    soup = BeautifulSoup(response.content, 'html.parser')
    
    leaderboard_div = soup.find('div', class_='flex flex-col gap-4 sticky top-20')
    
    if not leaderboard_div:
        raise ValueError("Leaderboard div not found. The page structure might have changed.")
    
    table = leaderboard_div.find('table', class_='w-full caption-bottom text-sm')
    
    if not table:
        raise ValueError("Leaderboard table not found within the div.")
    
    data = []
    for row in table.find('tbody').find_all('tr'):
        cols = row.find_all('td')
        rank = cols[0].find('div', class_='flex').text.strip().split()[0]
        model = cols[0].find('a').text.strip()
        score = cols[1].text.strip()
        confidence = cols[2].text.strip()
        data.append([rank, model, score, confidence])
    
    df = pd.DataFrame(data, columns=['Rank', 'Model', 'Score', '95% Confidence'])
    return df

def update_leaderboard(leaderboard):
    try:
        df = scrape_leaderboard(leaderboard)
        return df, df.to_html(index=False)
    except Exception as e:
        return None, f"An error occurred: {str(e)}"

def export_to_excel(df):
    if df is not None:
        output = io.BytesIO()
        with pd.ExcelWriter(output, engine='openpyxl') as writer:
            df.to_excel(writer, index=False, sheet_name='Leaderboard')
        output.seek(0)
        return output
    return None

# Create Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# Scale AI Leaderboard Viewer")
    with gr.Row():
        dropdown = gr.Dropdown(choices=list(LEADERBOARDS.keys()), label="Select Leaderboard", value="Coding")
        export_button = gr.Button("Export to Excel")
    
    table_output = gr.HTML()
    df_state = gr.State()
    
    def on_load():
        df, html = update_leaderboard("Coding")
        return df, html
    
    dropdown.change(update_leaderboard, inputs=[dropdown], outputs=[df_state, table_output])
    export_button.click(export_to_excel, inputs=[df_state], outputs=[gr.File(label="Download Excel")])
    
    iface.load(on_load, outputs=[df_state, table_output])

# Launch the app
iface.launch()