#!/usr/bin/env python3
"""Build viewer data for the Oolong viewer.

Oolong (https://github.com/abertsch72/oolong, arXiv:2511.02817) is a long-context
reasoning/aggregation benchmark with two datasets:

  synth  — oolongbench/oolong-synth (parquet). Fields per row: id, context_len,
           dataset (source), context_window_text[_with_labels], question,
           task_group, task, answer, answer_type, num_labels, context_window_id.
  real   — oolongbench/oolong-real, toy_dnd config (jsonl, test=campaign1 +
           validation=campaign2). Fields: id, context_window_id,
           context_window_text, question, answer, question_type, episodes, campaign.

In both, many questions share one long ``context_window_text`` (keyed by
``context_window_id``) — that is the natural corpus. This script dedups the
context windows into per-document shards (capped, see CAP) and projects the
questions into a compact eval file. It writes, per set:

  corpus_<set>.json   list[{cwid, title, size, truncated, n_questions, file, meta...}]
  corpus_<set>/*.txt  one (capped) context-window shard, lazy-loaded
  eval_<set>.json     list[{id, context_window_id, question, answer, meta{...}}]

plus a ``sets.json`` manifest with per-set counts + filter facets.

Run from the viewer repo root:

    python scripts/build_data.py --data-dir /mnt/tmp/oolong
"""
import argparse
import ast
import json
import os
import re
import shutil

try:
    import pyarrow.parquet as pq
except ImportError:
    pq = None

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_DATA_DIR = "/mnt/tmp/oolong"
CAP = 2_000_000   # max chars kept per context-window shard (long-context benchmark)


def slugify(idx, s):
    safe = re.sub(r"[^0-9A-Za-z._-]+", "_", str(s)).strip("_")[:60] or "cw"
    return f"{idx:04d}_{safe}"


def clean_answer(a):
    if a is None:
        return ""
    s = str(a).strip()
    m = re.search(r"datetime\.date\((\d+),\s*(\d+),\s*(\d+)\)", s)
    if m:
        y, mo, d = (int(x) for x in m.groups())
        return f"{y:04d}-{mo:02d}-{d:02d}"
    try:
        v = ast.literal_eval(s)
        if isinstance(v, (list, tuple)):
            return ", ".join(str(x) for x in v)
        return str(v)
    except Exception:
        return s.strip("[]").strip()


def write_shard(shard_dir, idx, cwid, text):
    truncated = len(text) > CAP
    body = text[:CAP]
    if truncated:
        body += ("\n\n… [truncated for the viewer at %d characters; "
                 "see the full context on Hugging Face] …" % CAP)
    fname = slugify(idx, cwid) + ".txt"
    with open(os.path.join(shard_dir, fname), "w", encoding="utf-8") as f:
        f.write(body)
    return fname, len(text), truncated


def emit(set_name, corpus, eval_rows, facets):
    corpus_rows = sorted(corpus.values(), key=lambda d: d["title"].lower())
    corpus_path = os.path.join(ROOT, f"corpus_{set_name}.json")
    eval_path = os.path.join(ROOT, f"eval_{set_name}.json")
    with open(corpus_path, "w", encoding="utf-8") as f:
        json.dump(corpus_rows, f, ensure_ascii=False)
    with open(eval_path, "w", encoding="utf-8") as f:
        json.dump(eval_rows, f, ensure_ascii=False)
    total_mb = sum(r["size"] for r in corpus_rows) / 1e6
    print(f"[{set_name}] questions={len(eval_rows)} contexts={len(corpus_rows)} "
          f"corpus_text~{total_mb:.1f}MB (capped {CAP/1e6:.0f}MB/shard) "
          f"index={os.path.getsize(corpus_path)/1e6:.2f}MB eval={os.path.getsize(eval_path)/1e6:.2f}MB")
    return {
        "set": set_name,
        "n_questions": len(eval_rows),
        "n_contexts": len(corpus_rows),
        "corpus_file": f"corpus_{set_name}.json",
        "eval_file": f"eval_{set_name}.json",
        "facets": facets,
    }


# ----------------------- synth (parquet) -----------------------
def build_synth(data_dir):
    if pq is None:
        raise SystemExit("pyarrow required for the synth (parquet) dataset")
    import glob
    shards = sorted(glob.glob(os.path.join(data_dir, "oolong-synth/data/*.parquet")))
    if not shards:
        print("[synth] no parquet shards found — skipping")
        return None
    shard_dir = os.path.join(ROOT, "corpus_synth")
    if os.path.isdir(shard_dir):
        shutil.rmtree(shard_dir)
    os.makedirs(shard_dir)

    corpus = {}          # cwid -> index row
    eval_rows = []
    facet_vals = {"Task group": set(), "Answer type": set(), "Source dataset": set()}
    cols = ["id", "context_len", "dataset", "context_window_text", "question",
            "task_group", "task", "answer", "answer_type", "num_labels",
            "context_window_id"]
    for sh in shards:
        avail = pq.read_schema(sh).names
        use = [c for c in cols if c in avail]
        d = pq.read_table(sh, columns=use).to_pydict()
        n = len(d["context_window_id"])
        for i in range(n):
            cwid = d["context_window_id"][i]
            src = d.get("dataset", [None] * n)[i]
            clen = d.get("context_len", [None] * n)[i]
            if cwid not in corpus:
                fname, size, trunc = write_shard(shard_dir, len(corpus), cwid,
                                                 d["context_window_text"][i] or "")
                corpus[cwid] = {
                    "cwid": str(cwid),
                    "title": f"{src} · {clen} tok · #{cwid}",
                    "source": src, "context_len": clen,
                    "size": size, "truncated": trunc, "n_questions": 0,
                    "file": f"corpus_synth/{fname}",
                }
            corpus[cwid]["n_questions"] += 1
            tg = d.get("task_group", [""] * n)[i]
            at = (d.get("answer_type", [""] * n)[i] or "").replace("ANSWER_TYPE.", "")
            tk = (d.get("task", [""] * n)[i] or "").replace("TASK_TYPE.", "")
            facet_vals["Task group"].add(tg)
            facet_vals["Answer type"].add(at)
            facet_vals["Source dataset"].add(src)
            eval_rows.append({
                "id": str(d["id"][i]),
                "context_window_id": str(cwid),
                "question": (d.get("question", [""] * n)[i] or "").strip(),
                "answer": clean_answer(d.get("answer", [""] * n)[i]),
                "meta": {
                    "Task group": tg, "Task": tk, "Answer type": at,
                    "Context length": str(clen), "Source dataset": src,
                    "# labels": str(d.get("num_labels", [""] * n)[i]),
                },
            })
    eval_rows.sort(key=lambda r: (r["meta"]["Source dataset"] or "", r["id"]))
    facets = [{"key": k, "values": sorted(str(v) for v in vals if v not in (None, ""))}
              for k, vals in facet_vals.items()]
    return emit("synth", corpus, eval_rows, facets)


# ----------------------- real (toy_dnd jsonl) -----------------------
def build_real(data_dir):
    files = [("test", "campaign1"), ("validation", "campaign2")]
    paths = [(sp, os.path.join(data_dir, f"oolong-real/toy_dnd/{sp}.jsonl")) for sp, _ in files]
    if not all(os.path.exists(p) for _, p in paths):
        print("[real] toy_dnd jsonl not found — skipping")
        return None
    shard_dir = os.path.join(ROOT, "corpus_real")
    if os.path.isdir(shard_dir):
        shutil.rmtree(shard_dir)
    os.makedirs(shard_dir)

    corpus = {}
    eval_rows = []
    facet_vals = {"Question type": set(), "Split": set(), "Campaign": set()}
    for split, path in paths:
        with open(path, encoding="utf-8") as f:
            for line in f:
                if not line.strip():
                    continue
                r = json.loads(line)
                cwid = r["context_window_id"]
                camp = r.get("campaign", "")
                eps = r.get("episodes")
                if cwid not in corpus:
                    fname, size, trunc = write_shard(shard_dir, len(corpus), cwid,
                                                     r.get("context_window_text") or "")
                    corpus[cwid] = {
                        "cwid": cwid,
                        "title": f"{camp} · ep {eps} · {str(cwid)[:8]}",
                        "campaign": camp, "episodes": eps,
                        "size": size, "truncated": trunc, "n_questions": 0,
                        "file": f"corpus_real/{fname}",
                    }
                corpus[cwid]["n_questions"] += 1
                qt = r.get("question_type", "")
                facet_vals["Question type"].add(qt)
                facet_vals["Split"].add(split)
                facet_vals["Campaign"].add(camp)
                eval_rows.append({
                    "id": str(r.get("id", "")),
                    "context_window_id": cwid,
                    "question": (r.get("question") or "").strip(),
                    "answer": clean_answer(r.get("answer")),
                    "meta": {
                        "Question type": qt, "Split": split, "Campaign": camp,
                        "Episodes": str(eps),
                    },
                })
    eval_rows.sort(key=lambda r: (r["meta"]["Split"], r["id"]))
    facets = [{"key": k, "values": sorted(str(v) for v in vals if v not in (None, ""))}
              for k, vals in facet_vals.items()]
    return emit("real", corpus, eval_rows, facets)


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--data-dir", default=DEFAULT_DATA_DIR,
                    help="dir with oolong-synth/ and oolong-real/ downloads")
    ap.add_argument("--only", choices=["synth", "real"], help="build only one set")
    args = ap.parse_args()
    manifest = []
    if args.only in (None, "synth"):
        m = build_synth(args.data_dir)
        if m:
            manifest.append(m)
    if args.only in (None, "real"):
        m = build_real(args.data_dir)
        if m:
            manifest.append(m)
    # merge with any existing manifest entries not rebuilt this run
    path = os.path.join(ROOT, "sets.json")
    existing = []
    if os.path.exists(path) and args.only:
        existing = [s for s in json.load(open(path)) if s["set"] != args.only]
    sets = existing + manifest
    order = {"synth": 0, "real": 1}
    sets.sort(key=lambda s: order.get(s["set"], 9))
    with open(path, "w", encoding="utf-8") as f:
        json.dump(sets, f, ensure_ascii=False, indent=2)
    print("wrote sets.json:", [s["set"] for s in sets])


if __name__ == "__main__":
    main()