"""Export saved benchmark reports into the experiment-tracker CSV format.

The tracker (``docs/experiment_tracker.csv``) is the planning checklist — one row
per planned run with the config filled in and the metric cells blank. This module
reads the reports saved by the Benchmark tab (stored on the HF dataset repo) and
emits rows in the *exact same column order*, with the metric cells populated, so
results can be collected in the format defined by the experiment plan.

Usage (from the repo root)::

    # Dump every saved report into docs/experiment_results.csv (overwrite):
    python -m app.export_tracker --all

    # Append a single report by its id:
    python -m app.export_tracker <report_id>

    # Tag the row with its planned experiment id / hypothesis note:
    python -m app.export_tracker <report_id> --exp T2.1 --phase B --note "BM25 wins on legal"

Config values are normalised to the tracker's short vocabulary (e.g.
``Alibaba-NLP/gte-modernbert-base`` -> ``GTE-ModernBERT``) so exported rows line
up with the planned rows in ``experiment_tracker.csv``.
"""

import argparse
import csv
import json
import os
from pathlib import Path

from app import reports

_ROOT = Path(__file__).resolve().parent.parent
_INDEX_CONFIG_PATH = _ROOT / "configs" / "index_config.json"
_RESULTS_CSV = _ROOT / "docs" / "experiment_results.csv"

# Column order — must stay identical to docs/experiment_tracker.csv
COLUMNS = [
    "exp_id", "phase", "dataset", "domain", "factor_changed", "factor_value",
    "embedding_model", "chunking", "retrieval", "query_transform", "reranker",
    "top_n", "fusion_top_k", "llm_model", "eval_method", "n_samples",
    "pred_relevance", "pred_adherence", "pred_completeness", "pred_utilization",
    "gold_relevance", "gold_adherence", "gold_completeness", "gold_utilization",
    "relevance_rmse", "utilization_rmse", "completeness_rmse", "hallucination_auroc",
    "abstention_pct", "latency_per_q_s", "run_seconds",
    "hypothesis_confirmed", "report_id", "notes",
]

_EMBED_SHORT = {
    "sentence-transformers/all-MiniLM-L6-v2":  "MiniLM-L6",
    "Alibaba-NLP/gte-modernbert-base":         "GTE-ModernBERT",
    "BAAI/bge-base-en-v1.5":                   "BGE-base",
    "BAAI/bge-large-en-v1.5":                  "BGE-large",
    "sentence-transformers/all-mpnet-base-v2": "MPNet",
    "nlpaueb/legal-bert-base-uncased":         "Legal-BERT",
}

_LLM_SHORT = {
    "llama-3.1-8b-instant":    "llama-3.1-8b",
    "llama-3.3-70b-versatile": "llama-3.3-70b",
    "openai/gpt-4o-mini":      "gpt-4o-mini",
    "openai/gpt-4o":           "gpt-4o",
}

# stats-table metric label -> tracker column
_STATS_MAP = {
    "Relevance RMSE":      "relevance_rmse",
    "Utilization RMSE":    "utilization_rmse",
    "Completeness RMSE":   "completeness_rmse",
    "Hallucination AUROC":  "hallucination_auroc",
    "Abstention Rate (%)":  "abstention_pct",
    "Reject Rate (%)":      "abstention_pct",  # legacy reports (pre-abstention rename)
    "Latency/query (s)":    "latency_per_q_s",
    "Run time (s)":         "run_seconds",
}


def _load_index_config() -> dict:
    with open(_INDEX_CONFIG_PATH) as f:
        return json.load(f)


def _norm_chunking(chunk_label: str) -> str:
    """'512/96 (recursive)' -> '512/96 recursive'; '0/0 (passthrough)' -> 'passthrough'."""
    if not chunk_label:
        return ""
    label = chunk_label.replace("(", "").replace(")", "").strip()
    if "passthrough" in label:
        return "passthrough"
    return label


def report_to_row(report: dict, index_config: dict,
                  exp_id: str = "", phase: str = "",
                  factor_changed: str = "", factor_value: str = "",
                  hypothesis_confirmed: str = "", note: str = "") -> dict:
    """Convert a saved report dict into a tracker-format row dict."""
    cfg = report.get("config", {})
    dataset = cfg.get("dataset", "")
    domain = (
        index_config.get("datasets", {}).get(dataset, {}).get("domain", "")
    )

    reranker_raw = cfg.get("reranker", "")
    row = {c: "" for c in COLUMNS}
    row.update({
        "exp_id":          exp_id,
        "phase":           phase,
        "dataset":         dataset,
        "domain":          domain,
        "factor_changed":  factor_changed,
        "factor_value":    factor_value,
        "embedding_model": _EMBED_SHORT.get(cfg.get("embedding_model", ""), cfg.get("embedding_model", "")),
        "chunking":        _norm_chunking(cfg.get("chunk_label", "")),
        "retrieval":       cfg.get("retrieval_strategy", ""),
        "query_transform": cfg.get("query_rewrite", ""),
        "reranker":        "off" if reranker_raw == "off" else "on",
        "top_n":           cfg.get("top_n", ""),
        "fusion_top_k":    cfg.get("fusion_top_k", ""),
        "llm_model":       _LLM_SHORT.get(cfg.get("llm_model", ""), cfg.get("llm_model", "")),
        "eval_method":     cfg.get("eval_method", ""),
        "n_samples":       cfg.get("n_samples", ""),
        "hypothesis_confirmed": hypothesis_confirmed,
        "report_id":       report.get("report_id", ""),
        "notes":           note,
    })

    # Predicted / gold means from the summary block
    for s in report.get("summary", []):
        m = s.get("metric", "")
        if m in ("relevance", "adherence", "completeness", "utilization"):
            row[f"pred_{m}"] = _fmt(s.get("pred_mean"))
            row[f"gold_{m}"] = _fmt(s.get("gold_mean"))

    # RMSE / AUROC / run diagnostics from the stats block
    for s in report.get("stats", []):
        col = _STATS_MAP.get(s.get("metric", ""))
        if col:
            row[col] = _fmt(s.get("value"))

    return row


def _fmt(v):
    if v is None:
        return ""
    try:
        f = float(v)
    except (TypeError, ValueError):
        return v
    if f != f:  # NaN
        return ""
    return round(f, 4)


def _write_rows(rows: list[dict], path: Path, append: bool) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    file_exists = path.exists()
    mode = "a" if append and file_exists else "w"
    with open(path, mode, newline="") as f:
        writer = csv.DictWriter(f, fieldnames=COLUMNS)
        if mode == "w":
            writer.writeheader()
        for r in rows:
            writer.writerow({c: r.get(c, "") for c in COLUMNS})


def main() -> None:
    parser = argparse.ArgumentParser(description="Export benchmark reports to tracker-format CSV.")
    parser.add_argument("report_id", nargs="?", help="Report id to append (omit with --all).")
    parser.add_argument("--all", action="store_true", help="Dump every saved report (overwrites the results CSV).")
    parser.add_argument("--exp", default="", help="Planned experiment id, e.g. T2.1.")
    parser.add_argument("--phase", default="", help="Phase label, e.g. B.")
    parser.add_argument("--factor", default="", help="factor_changed value.")
    parser.add_argument("--value", default="", help="factor_value value.")
    parser.add_argument("--confirmed", default="", help="hypothesis_confirmed value (yes/no/partial).")
    parser.add_argument("--note", default="", help="Free-text note.")
    parser.add_argument("--out", default=str(_RESULTS_CSV), help="Output CSV path.")
    args = parser.parse_args()

    index_config = _load_index_config()
    hub_repo = index_config.get("hub_repo", "")
    token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN")
    out_path = Path(args.out)

    if args.all:
        idx = reports.load_reports_index(hub_repo, token=token)
        if not idx:
            print("No reports found in the index.")
            return
        rows = []
        for entry in idx:
            rep = reports.load_report(entry["report_id"], hub_repo, token=token)
            if rep:
                rows.append(report_to_row(rep, index_config))
        rows.sort(key=lambda r: (r["dataset"], r["report_id"]))
        _write_rows(rows, out_path, append=False)
        print(f"Wrote {len(rows)} report row(s) to {out_path}")
        return

    if not args.report_id:
        parser.error("provide a report_id or use --all")

    rep = reports.load_report(args.report_id, hub_repo, token=token)
    if not rep:
        print(f"Report {args.report_id} could not be loaded.")
        return
    row = report_to_row(
        rep, index_config,
        exp_id=args.exp, phase=args.phase,
        factor_changed=args.factor, factor_value=args.value,
        hypothesis_confirmed=args.confirmed, note=args.note,
    )
    _write_rows([row], out_path, append=True)
    print(f"Appended report {args.report_id} to {out_path}")


if __name__ == "__main__":
    main()