"""Disk cache for Benchmark-tab results, keyed by (config_hash, row_index).

Re-running the Benchmark tab with the same config + sample reuses cached
scores instead of re-running retrieval/generation/API calls, mirroring the
checkpointing pattern from Notebooks 5-6.

Gold scores (from RAGBench labels) are stored alongside predicted scores so
that RMSE/AUROC can be recomputed from cache without re-querying the dataset.
"""

import hashlib
import json

import pandas as pd

from src import settings

RESULTS_PATH = settings.APP_CACHE_DIR / "benchmark_results.csv"
COLUMNS = [
    "config_hash", "row_index", "question",
    # predicted scores
    "relevance", "adherence", "completeness", "utilization",
    # RAGBench gold labels
    "gold_adherence", "gold_relevance", "gold_utilization", "gold_completeness",
    # run diagnostics
    "rejected", "abstained", "latency_s",
]


def config_hash(dataset: str, index_id: str, config_dict: dict) -> str:
    raw = json.dumps({"dataset": dataset, "index_id": index_id, **config_dict}, sort_keys=True)
    return hashlib.md5(raw.encode("utf-8")).hexdigest()[:12]


def load_results() -> pd.DataFrame:
    if RESULTS_PATH.exists():
        df = pd.read_csv(RESULTS_PATH)
        # Back-fill columns added after older cache files were written
        for col in ["gold_adherence", "gold_relevance", "gold_utilization", "gold_completeness",
                    "rejected", "abstained", "latency_s"]:
            if col not in df.columns:
                df[col] = float("nan")
        return df
    return pd.DataFrame(columns=COLUMNS)


def save_results(df: pd.DataFrame) -> None:
    df.to_csv(RESULTS_PATH, index=False)


def get_cached_rows(results_df: pd.DataFrame, chash: str) -> pd.DataFrame:
    return results_df[results_df["config_hash"] == chash]


def upsert_rows(results_df: pd.DataFrame, chash: str, new_rows: pd.DataFrame) -> pd.DataFrame:
    """Replace any rows for (chash, row_index) with new_rows, keeping everything else."""
    new_keys = set(zip(new_rows["row_index"], [chash] * len(new_rows)))
    keep_mask = ~results_df.apply(
        lambda r: (r["row_index"], r["config_hash"]) in new_keys, axis=1
    )
    merged = pd.concat([results_df[keep_mask], new_rows], ignore_index=True)
    save_results(merged)
    return merged