import gradio as gr
import pandas as pd
from huggingface_hub import HfFileSystem

RESULTS_REPO = "CEIA-COREJUR/llm-benchmark-br-resultados"


def load_results() -> pd.DataFrame:
    fs = HfFileSystem()
    dfs = []

    try:
        entries = fs.ls(f"datasets/{RESULTS_REPO}", detail=False)
    except Exception:
        return pd.DataFrame()

    for entry in entries:
        folder_name = entry.split("/")[-1]

        model_info = {}
        try:
            with fs.open(f"{entry}/model.txt", "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if ":" in line and not line.startswith("#"):
                        key, _, val = line.partition(":")
                        model_info[key.strip()] = val.strip()
        except Exception:
            pass

        model_label = model_info.get("variant_name", folder_name)
        backbone = model_info.get("backbone", "")
        provider = model_info.get("provider", "")

        try:
            csv_files = [f for f in fs.ls(entry, detail=False) if f.endswith(".csv")]
        except Exception:
            continue

        for csv_path in csv_files:
            exam_slug = csv_path.split("/")[-1].replace(".csv", "")
            try:
                with fs.open(csv_path, "r", encoding="utf-8") as f:
                    df = pd.read_csv(f, dtype=str)

                # Nome real da prova (exam_edition); CSVs antigos sem essa coluna caem
                # de volta pro nome do arquivo, só pra não quebrar o filtro.
                if "exam_edition" in df.columns and not df["exam_edition"].dropna().empty:
                    exam_name = df["exam_edition"].iloc[0]
                    df = df.drop(columns=["exam_edition"])
                else:
                    exam_name = exam_slug

                df.insert(0, "model", model_label)
                df.insert(1, "backbone", backbone)
                df.insert(2, "provider", provider)
                df.insert(3, "exam", exam_name)
                dfs.append(df)
            except Exception:
                continue

    if not dfs:
        return pd.DataFrame()

    combined = pd.concat(dfs, ignore_index=True)
    combined["is_correct"] = combined["is_correct"].map({"True": True, "False": False})
    return combined


def build_leaderboard(df: pd.DataFrame) -> pd.DataFrame:
    if df.empty:
        return pd.DataFrame()
    summary = (
        df.groupby(["model", "backbone", "provider"])
        .agg(total=("is_correct", "count"), correct=("is_correct", "sum"))
        .reset_index()
    )
    summary["accuracy"] = (summary["correct"] / summary["total"] * 100).round(1).astype(str) + "%"
    return summary[["model", "backbone", "provider", "accuracy", "correct", "total"]]


def filter_results(model, exam, subject, correct):
    df = df_global.copy()
    if model != "Todos":
        df = df[df["model"] == model]
    if exam != "Todos":
        df = df[df["exam"] == exam]
    if subject != "Todos":
        df = df[df["subject"] == subject]
    if correct == "Corretas":
        df = df[df["is_correct"] == True]
    elif correct == "Erradas":
        df = df[df["is_correct"] == False]
    return df


def choices(col):
    if df_global.empty or col not in df_global.columns:
        return ["Todos"]
    return ["Todos"] + sorted(df_global[col].dropna().unique().tolist())


df_global = load_results()

with gr.Blocks(title="Benchmark LLMs — Questões BR") as demo:
    gr.Markdown("# Benchmark LLMs — Questões de Múltipla Escolha em Português")

    with gr.Tab("Leaderboard"):
        gr.Dataframe(value=build_leaderboard(df_global), interactive=False)

    with gr.Tab("Resultados"):
        with gr.Row():
            model_filter   = gr.Dropdown(choices=choices("model"),   value="Todos", label="Modelo")
            exam_filter    = gr.Dropdown(choices=choices("exam"),    value="Todos", label="Prova")
            subject_filter = gr.Dropdown(choices=choices("subject"), value="Todos", label="Disciplina")
            correct_filter = gr.Dropdown(choices=["Todos", "Corretas", "Erradas"], value="Todos", label="Resultado")

        results_table = gr.Dataframe(value=df_global, interactive=False)

        for comp in [model_filter, exam_filter, subject_filter, correct_filter]:
            comp.change(
                filter_results,
                inputs=[model_filter, exam_filter, subject_filter, correct_filter],
                outputs=results_table,
            )

demo.launch()