#!/usr/bin/env python3
"""Build eval.json + products.json for the HERB viewer.

HERB ships one JSON per product, each bundling heterogeneous *artifacts*
(slack / documents / meeting_transcripts / meeting_chats / urls / prs) plus
answerable + unanswerable questions. The viewer streams the (heavy) artifact
content on demand straight from the Salesforce/HERB dataset on the HF CDN, so
this script only derives the small bundled files:

- eval.json      one entry per question (815 answerable + 699 unanswerable);
                 product, kind, question, ground_truth, citations (evidence
                 artifact ids), type, and computation/comparison label.
- products.json  one entry per product with per-type artifact counts and
                 question counts (drives the Corpus/Eval product filters).

employees.json / customers.json (id -> name maps) are copied verbatim from the
dataset's metadata/ folder and bundled for client-side id resolution.

Run from the viewer repo root:
    python scripts/build_data.py [--herb-dir DIR]

Reads:  <herb-dir>/products/*.json
Writes: eval.json, products.json
"""
import argparse
import glob
import json
import os
import re

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_HERB = "/mnt/ramdisk/blobstore/timchen0618/data/HERB"
ARTIFACT_TYPES = ["slack", "documents", "meeting_transcripts",
                  "meeting_chats", "urls", "prs"]

COMPUTATION_COMPARISON_PATTERNS = [
    (
        "maximum_or_minimum",
        re.compile(
            r"\b(?:maximum|minimum|highest|lowest|largest|smallest|"
            r"most|least|fewest|greatest)\b",
            re.IGNORECASE,
        ),
    ),
    (
        "numeric_comparison",
        re.compile(
            r"\b(?:compare|comparison|versus|vs\.?|difference between|"
            r"more than|less than|greater than|fewer than|higher than|"
            r"lower than)\b",
            re.IGNORECASE,
        ),
    ),
    (
        "arithmetic",
        re.compile(
            r"\b(?:sum|total|average|mean|percentage|percent|ratio|"
            r"how many|combined|altogether|remaining)\b",
            re.IGNORECASE,
        ),
    ),
]


def computation_comparison(question):
    """Classify questions using the benchmark split requested for analysis."""
    for reason, pattern in COMPUTATION_COMPARISON_PATTERNS:
        if pattern.search(question or ""):
            return True, reason
    return False, "none"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--herb-dir", default=DEFAULT_HERB)
    args = ap.parse_args()

    prod_files = sorted(glob.glob(os.path.join(args.herb_dir, "products", "*.json")))
    if not prod_files:
        raise SystemExit(f"no product json under {args.herb_dir}/products/")

    eval_rows = []
    products = []
    for pf in prod_files:
        name = os.path.basename(pf)[:-5]
        d = json.load(open(pf, encoding="utf-8"))

        counts = {t: len(d.get(t, []) or []) for t in ARTIFACT_TYPES}
        aq = d.get("answerable_questions", []) or []
        uq = d.get("unanswerable_questions", []) or []
        products.append({
            "name": name,
            "counts": counts,
            "n_artifacts": sum(counts.values()),
            "n_answerable": len(aq),
            "n_unanswerable": len(uq),
        })

        for i, q in enumerate(aq):
            is_computation_comparison, category_reason = computation_comparison(
                q.get("question", "")
            )
            eval_rows.append({
                "gid": f"{name}#a{i}",
                "product": name,
                "kind": "answerable",
                "question": q.get("question", ""),
                "ground_truth": q.get("ground_truth"),
                "citations": q.get("citations", []) or [],
                "type": q.get("type", ""),
                "computation_comparison": is_computation_comparison,
                "computation_comparison_reason": category_reason,
            })
        for i, q in enumerate(uq):
            # unanswerable questions are plain strings
            qtext = q if isinstance(q, str) else (q.get("question", "") if isinstance(q, dict) else str(q))
            is_computation_comparison, category_reason = computation_comparison(qtext)
            eval_rows.append({
                "gid": f"{name}#u{i}",
                "product": name,
                "kind": "unanswerable",
                "question": qtext,
                "ground_truth": None,
                "citations": [],
                "type": "unanswerable",
                "computation_comparison": is_computation_comparison,
                "computation_comparison_reason": category_reason,
            })

    with open(os.path.join(ROOT, "eval.json"), "w", encoding="utf-8") as f:
        json.dump(eval_rows, f, ensure_ascii=False, indent=0)
    with open(os.path.join(ROOT, "products.json"), "w", encoding="utf-8") as f:
        json.dump(products, f, ensure_ascii=False, indent=2)

    tot_art = sum(p["n_artifacts"] for p in products)
    n_ans = sum(p["n_answerable"] for p in products)
    n_unans = sum(p["n_unanswerable"] for p in products)
    print(f"wrote eval.json: {len(eval_rows)} questions ({n_ans} answerable, {n_unans} unanswerable)")
    print(f"wrote products.json: {len(products)} products, {tot_art} artifacts total")
    from collections import Counter
    print("answerable types:", dict(Counter(r["type"] for r in eval_rows if r["kind"] == "answerable")))
    for kind in ("answerable", "unanswerable"):
        bracket_counts = Counter(
            row["computation_comparison"]
            for row in eval_rows
            if row["kind"] == kind
        )
        print(
            f"{kind} computation/comparison: "
            f"{bracket_counts[True]} yes, {bracket_counts[False]} no"
        )


if __name__ == "__main__":
    main()