#!/usr/bin/env python3
"""Export BFSI-Bench preview rows to data/samples.json for the static Space.

The Space is public, so this writes at most PER_CATEGORY items per
category_code (sorted by id). The full test set stays on the gated Hub
dataset.

Prefers local hf_dataset/bfsi-bench/data/test.csv when present; otherwise
downloads the Hub dataset ground-truth/bfsi-bench.
"""

from __future__ import annotations

import csv
import io
import json
from pathlib import Path

_SPACE = Path(__file__).resolve().parents[1]
_REPO = Path(__file__).resolve().parents[3]
_LOCAL_CSV = _REPO / "hf_dataset" / "bfsi-bench" / "data" / "test.csv"
_OUT = _SPACE / "data" / "samples.json"

HUB_ID = "ground-truth/bfsi-bench"
JSON_COLS = ("sources", "evidence", "gold_nuggets")
PER_CATEGORY = 5


def _parse_json_col(raw: str, *, field: str, row_id: str):
    text = (raw or "").strip()
    if not text:
        return []
    try:
        val = json.loads(text)
    except json.JSONDecodeError as exc:
        raise SystemExit(f"{row_id}: invalid JSON in {field}: {exc}") from exc
    if val is None:
        return []
    if not isinstance(val, list):
        raise SystemExit(f"{row_id}: {field} must be a JSON list")
    return val


def _rows_from_csv(text: str) -> list[dict]:
    reader = csv.DictReader(io.StringIO(text))
    items = []
    for raw in reader:
        rid = (raw.get("id") or "").strip()
        if not rid:
            continue
        item = {
            "id": rid,
            "category": (raw.get("category") or "").strip(),
            "category_code": (raw.get("category_code") or "").strip(),
            "query": raw.get("query") or "",
            "gold_answer": raw.get("gold_answer") or "",
        }
        for col in JSON_COLS:
            item[col] = _parse_json_col(raw.get(col) or "", field=col, row_id=rid)
        items.append(item)
    items.sort(key=lambda r: r["id"])
    return items


def _counts_by_category(items: list[dict]) -> dict[str, int]:
    counts: dict[str, int] = {}
    for it in items:
        code = it["category_code"] or "unknown"
        counts[code] = counts.get(code, 0) + 1
    return counts


def take_per_category(
    items: list[dict], n: int = PER_CATEGORY
) -> list[dict]:
    """Keep the first n items in each category_code. Items must already be sorted by id."""
    by_cat: dict[str, list[dict]] = {}
    for it in items:
        code = it["category_code"] or "unknown"
        bucket = by_cat.setdefault(code, [])
        if len(bucket) < n:
            bucket.append(it)
    sampled: list[dict] = []
    for code in sorted(by_cat):
        sampled.extend(by_cat[code])
    sampled.sort(key=lambda r: r["id"])
    return sampled


def _load_csv_text() -> tuple[str, str]:
    if _LOCAL_CSV.is_file():
        return _LOCAL_CSV.read_text(encoding="utf-8-sig"), str(_LOCAL_CSV)
    try:
        from huggingface_hub import hf_hub_download
    except ImportError:
        raise SystemExit(
            "Local test.csv missing and huggingface_hub is not installed. "
            "pip install huggingface_hub or place hf_dataset/bfsi-bench/data/test.csv"
        ) from None
    path = hf_hub_download(
        repo_id=HUB_ID,
        repo_type="dataset",
        filename="data/test.csv",
    )
    return Path(path).read_text(encoding="utf-8-sig"), f"hub:{HUB_ID}/data/test.csv"


def main() -> None:
    csv_text, source = _load_csv_text()
    items = _rows_from_csv(csv_text)
    dataset_counts = _counts_by_category(items)
    preview = take_per_category(items, PER_CATEGORY)
    preview_counts = _counts_by_category(preview)
    payload = {
        "dataset": HUB_ID,
        "split": "test",
        "preview_per_category": PER_CATEGORY,
        "dataset_total": len(items),
        "dataset_counts": dataset_counts,
        "total": len(preview),
        "counts": preview_counts,
        "items": preview,
    }
    _OUT.parent.mkdir(parents=True, exist_ok=True)
    _OUT.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(
        json.dumps(
            {
                "wrote": str(_OUT.relative_to(_REPO)),
                "source": source,
                "preview_per_category": PER_CATEGORY,
                "dataset_total": len(items),
                "total": len(preview),
                "counts": preview_counts,
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()