#!/usr/bin/env python3
"""Export the four SLAVA leaderboard tables from an XLSX workbook to CSV.

The workbook uses ``model_name`` as the unique model identifier on the
Leaderboard sheet and ``model`` as the same identifier on the three slice
sheets.  The human-readable name on Leaderboard is preserved in
``display_name``.  Keeping those concepts separate prevents duplicate display
names (for example, Instruct and Think variants) from corrupting metadata
joins in the Streamlit application.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import pandas as pd


SHEET_TO_CSV = {
    "Leaderboard": "leaderboard.csv",
    "Subject": "subject.csv",
    "Type of question": "question_type.csv",
    "Provocativeness": "provocativeness.csv",
}

REQUIRED_LEADERBOARD_COLUMNS = {
    "model_name",
    "model",
    "RU_adaptation",
    "Access",
    "Team",
    "Country",
    "n_params, B",
    "rank",
    "LEADERBOARD",
}


def _clean_text_column(series: pd.Series) -> pd.Series:
    return series.astype("string").str.strip()


def _load_tables(workbook: Path) -> dict[str, pd.DataFrame]:
    excel = pd.ExcelFile(workbook)
    missing_sheets = set(SHEET_TO_CSV) - set(excel.sheet_names)
    if missing_sheets:
        raise ValueError(f"Missing workbook sheets: {sorted(missing_sheets)}")

    tables = {
        sheet: pd.read_excel(workbook, sheet_name=sheet).dropna(how="all")
        for sheet in SHEET_TO_CSV
    }

    leaderboard = tables["Leaderboard"]
    missing_columns = REQUIRED_LEADERBOARD_COLUMNS - set(leaderboard.columns)
    if missing_columns:
        raise ValueError(f"Leaderboard is missing columns: {sorted(missing_columns)}")

    leaderboard = leaderboard.rename(columns={"model": "display_name", "model_name": "model"})
    leaderboard["model"] = _clean_text_column(leaderboard["model"])
    leaderboard["display_name"] = _clean_text_column(leaderboard["display_name"])
    tables["Leaderboard"] = leaderboard

    for sheet in ("Subject", "Type of question", "Provocativeness"):
        table = tables[sheet]
        if "model" not in table.columns:
            raise ValueError(f"{sheet} is missing the model column")
        table["model"] = _clean_text_column(table["model"])
        tables[sheet] = table

    return tables


def _validate_tables(tables: dict[str, pd.DataFrame], expected_models: int) -> None:
    leaderboard = tables["Leaderboard"]
    if len(leaderboard) != expected_models:
        raise ValueError(
            f"Leaderboard has {len(leaderboard)} rows; expected {expected_models}"
        )
    if leaderboard["model"].isna().any() or leaderboard["model"].eq("").any():
        raise ValueError("Leaderboard contains an empty model identifier")
    if leaderboard["model"].duplicated().any():
        duplicates = leaderboard.loc[leaderboard["model"].duplicated(False), "model"].tolist()
        raise ValueError(f"Leaderboard contains duplicate model identifiers: {duplicates}")
    if leaderboard["rank"].duplicated().any():
        raise ValueError("Leaderboard contains duplicate ranks")

    expected_ids = set(leaderboard["model"])
    for sheet, table in tables.items():
        if len(table) != expected_models:
            raise ValueError(f"{sheet} has {len(table)} rows; expected {expected_models}")
        if table["model"].duplicated().any():
            raise ValueError(f"{sheet} contains duplicate model identifiers")
        actual_ids = set(table["model"])
        if actual_ids != expected_ids:
            missing = sorted(expected_ids - actual_ids)
            extra = sorted(actual_ids - expected_ids)
            raise ValueError(f"{sheet} model mismatch; missing={missing}, extra={extra}")


def export_workbook(workbook: Path, output_dir: Path, expected_models: int = 35) -> None:
    tables = _load_tables(workbook)
    _validate_tables(tables, expected_models)
    output_dir.mkdir(parents=True, exist_ok=True)

    for sheet, filename in SHEET_TO_CSV.items():
        destination = output_dir / filename
        tables[sheet].to_csv(destination, index=False, encoding="utf-8")
        print(f"{sheet}: {len(tables[sheet])} models -> {destination}")


def parse_args() -> argparse.Namespace:
    project_root = Path(__file__).resolve().parent.parent
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "workbook",
        type=Path,
        nargs="?",
        default=project_root / "source/slava_leaderboard_2026_q4.xlsx",
        help="Source leaderboard XLSX",
    )
    parser.add_argument("--output-dir", type=Path, default=project_root / "data")
    parser.add_argument("--expected-models", type=int, default=35)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    export_workbook(args.workbook, args.output_dir, args.expected_models)


if __name__ == "__main__":
    main()