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

MEBench (https://github.com/tl2309/MEBench, HF dataset Tim999999/MEBench) is a
cross-document multi-entity QA benchmark. It ships **questions only** — for each
of 3 splits (train / test / single_test) a JSONL of
``{qid, topic, edge, properties, type, question, answer, ...}`` over 7 "topics"
(university associations: Ivy League, Group of Eight, ...). The actual *corpus*
in MEBench is each entity's Wikipedia intro paragraph, generated live by the
pipeline and **not shipped**.

This script therefore reconstructs the corpus from the entities that appear in
the dataset (the ``Entity`` field of the test + single_test splits, which
together cover all 7 topics), fetching each entity's Wikipedia intro. It writes:

  corpus.json         list[{title, topics, size, wiki_url, file}]   (shared, split-independent)
  corpus/<slug>.txt   one Wikipedia-intro shard per unique entity (lazy-loaded)
  eval_<split>.json   the split's questions, projected + supporting_titles resolved
  sets.json           manifest of the 3 splits (counts, types, file pointers)

Wikipedia intros are cached to ``<data_dir>/wiki_cache.jsonl`` so re-runs are cheap.

Run from the viewer repo root:

    python scripts/build_data.py \
        --data-dir /mnt/ramdisk/blobstore/timchen0618/data/mebench
"""
import argparse
import json
import os
import re
import time
import shutil

try:
    import requests
except ImportError:
    requests = None

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_DATA_DIR = "/mnt/ramdisk/blobstore/timchen0618/data/mebench"
SPLITS = ["train", "test", "single_test"]

# type (normalized) -> reasoning category (paper: 3 categories, 8 types)
TYPE_CATEGORY = {
    "Intercomparison": "Comparative Reasoning",
    "Superlative": "Comparative Reasoning",
    "Aggregation": "Statistical Reasoning",
    "Distribution Compliance": "Statistical Reasoning",
    "Correlation Analysis": "Statistical Reasoning",
    "Variance Analysis": "Statistical Reasoning",
    "Descriptive Relationship": "Relational Reasoning",
    "Hypothetical Scenarios": "Relational Reasoning",
}


def norm_type(t):
    t = (t or "").strip()
    if t == "Hypthetical Scenarios":   # dataset typo
        return "Hypothetical Scenarios"
    return t


def norm_entity(e):
    return (e or "").strip().replace("_", " ").strip()


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


def load_jsonl(path):
    with open(path, encoding="utf-8") as f:
        return [json.loads(l) for l in f if l.strip()]


# ----------------------- Wikipedia corpus -----------------------
def load_cache(cache_path):
    cache = {}
    if os.path.exists(cache_path):
        for r in load_jsonl(cache_path):
            cache[r["title"]] = r["extract"]
    return cache


def fetch_wikipedia_intros(titles, cache_path):
    """Return {title: intro_text}. Batched (20/req), cached, redirects resolved."""
    cache = load_cache(cache_path)
    todo = [t for t in titles if t not in cache]
    if todo and requests is None:
        raise SystemExit("`requests` is required to fetch Wikipedia intros")
    session = requests.Session() if todo else None
    if session:
        session.headers.update({"User-Agent": "mebench-viewer/1.0 (dataset viewer; research)"})
    fout = open(cache_path, "a", encoding="utf-8") if todo else None
    for i in range(0, len(todo), 20):
        batch = todo[i:i + 20]
        params = {
            "action": "query", "prop": "extracts", "exintro": 1,
            "explaintext": 1, "exlimit": 20, "redirects": 1,
            "format": "json", "titles": "|".join(batch),
        }
        got = {}
        try:
            data = session.get("https://en.wikipedia.org/w/api.php",
                               params=params, timeout=60).json()
            # map redirected/normalized titles back to what we requested
            alias = {}
            for n in data.get("query", {}).get("normalized", []):
                alias[n["to"]] = n["from"]
            for rd in data.get("query", {}).get("redirects", []):
                alias[rd["to"]] = alias.get(rd["from"], rd["from"])
            for pg in data.get("query", {}).get("pages", {}).values():
                ret = pg.get("title", "")
                req = alias.get(ret, ret)
                got[req] = pg.get("extract", "") or ""
        except Exception as e:
            print("  wiki batch error:", str(e)[:120])
        for t in batch:
            ex = got.get(t, "")
            cache[t] = ex
            fout.write(json.dumps({"title": t, "extract": ex}, ensure_ascii=False) + "\n")
        fout.flush()
        time.sleep(0.2)
    if fout:
        fout.close()
    return {t: cache.get(t, "") for t in titles}


def build_corpus(data_dir):
    te = load_jsonl(os.path.join(data_dir, "test.jsonl"))
    si = load_jsonl(os.path.join(data_dir, "single_test.jsonl"))
    ent_topics = {}
    for d in te + si:
        e = norm_entity(d.get("Entity", ""))
        if not e:
            continue
        ent_topics.setdefault(e, set()).add(d["topic"].strip())

    titles = sorted(ent_topics)
    print(f"corpus: {len(titles)} unique entities across "
          f"{len({t for ts in ent_topics.values() for t in ts})} topics")
    intros = fetch_wikipedia_intros(titles, os.path.join(data_dir, "wiki_cache.jsonl"))

    shard_dir = os.path.join(ROOT, "corpus")
    if os.path.isdir(shard_dir):
        shutil.rmtree(shard_dir)
    os.makedirs(shard_dir)

    index_rows = []
    missing = 0
    for i, title in enumerate(titles):
        content = intros.get(title, "").strip()
        if not content:
            missing += 1
            content = "(No Wikipedia intro found for this entity.)"
        fname = slugify(i, title) + ".txt"
        with open(os.path.join(shard_dir, fname), "w", encoding="utf-8") as f:
            f.write(content)
        index_rows.append({
            "title": title,
            "topics": sorted(ent_topics[title]),
            "size": len(content),
            "wiki_url": "https://en.wikipedia.org/wiki/" + title.replace(" ", "_"),
            "file": f"corpus/{fname}",
        })
    index_rows.sort(key=lambda d: d["title"].lower())
    with open(os.path.join(ROOT, "corpus.json"), "w", encoding="utf-8") as f:
        json.dump(index_rows, f, ensure_ascii=False)
    print(f"corpus: wrote {len(index_rows)} shards "
          f"({missing} without a Wikipedia intro), "
          f"{sum(r['size'] for r in index_rows)/1e6:.2f}MB text")
    return {r["title"] for r in index_rows}, ent_topics


# ----------------------- eval splits -----------------------
def build_eval(data_dir, corpus_titles):
    manifest = []
    for split in SPLITS:
        rows = load_jsonl(os.path.join(data_dir, f"{split}.jsonl"))
        out = []
        types = {}
        for r in rows:
            topic = (r.get("topic") or "").strip()
            typ = norm_type(r.get("type"))
            types[typ] = types.get(typ, 0) + 1
            edge = r.get("edge")
            if isinstance(edge, list):
                edge = ", ".join(edge)
            entity = norm_entity(r.get("Entity", ""))
            # supporting docs: the specific entity if present & known. Otherwise
            # the question ranges over the whole topic (train is topic-level) —
            # we flag it and let the UI derive members from corpus.json (which
            # tags every entity with its topics) rather than duplicating the
            # (up to ~168-entry) member list on every question.
            if entity and entity in corpus_titles:
                supporting = [entity]
                topic_level = False
            else:
                supporting = []
                topic_level = True
            row = {
                "qid": r.get("qid"),
                "topic": topic,
                "type": typ,
                "category": TYPE_CATEGORY.get(typ, ""),
                "question": (r.get("question") or "").strip(),
                "answer": r.get("answer"),
                "edge": (edge or "").strip(),
                "properties": (r.get("properties") or "").strip(),
                "supporting_titles": supporting,
                "topic_level": topic_level,
            }
            if r.get("sql"):
                row["sql"] = r["sql"].strip()
            if r.get("class"):
                row["class"] = r["class"].strip()
            if r.get("Hops"):
                row["hops"] = r["Hops"].strip()
            if entity:
                row["entity"] = entity
            out.append(row)
        out.sort(key=lambda d: (d["topic"], d.get("qid") or 0))
        path = os.path.join(ROOT, f"eval_{split}.json")
        with open(path, "w", encoding="utf-8") as f:
            json.dump(out, f, ensure_ascii=False)
        manifest.append({
            "split": split,
            "n_questions": len(out),
            "topics": sorted({d["topic"] for d in out}),
            "types": dict(sorted(types.items(), key=lambda kv: -kv[1])),
            "eval_file": f"eval_{split}.json",
        })
        print(f"[{split}] questions={len(out)} "
              f"({os.path.getsize(path)/1e6:.2f}MB)")
    with open(os.path.join(ROOT, "sets.json"), "w", encoding="utf-8") as f:
        json.dump(manifest, f, ensure_ascii=False, indent=2)


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--data-dir", default=DEFAULT_DATA_DIR,
                    help="dir with train/test/single_test.jsonl (+ wiki_cache.jsonl)")
    args = ap.parse_args()
    corpus_titles, ent_topics = build_corpus(args.data_dir)
    build_eval(args.data_dir, corpus_titles)


if __name__ == "__main__":
    main()