#!/usr/bin/env python3
"""Build the FinQA viewer data from the raw dataset splits.

Reads the four raw split files (train/dev/test/private_test.json) and writes:

    eval.jsonl                 one slim record per QA example (9,200 lines)
    corpus/index.json          light per-document index (2,789 unique pages)
    corpus/records/<key>.json  full page content (pre_text + table + post_text)

A "document" is a unique source page identified by `filename`
(`TICKER/YEAR/page_N.pdf`). Page content is identical across all examples that
share a filename, so it is de-duplicated into the corpus.

Run from the viewer repo root:
    python scripts/build_data.py [--src <dataset-dir>]

Default --src is $DATA_ROOT/finqa/dataset (falls back to ./dataset).
"""
import argparse
import json
import os
from collections import Counter, defaultdict

SPLITS = ["train", "dev", "test", "private_test"]
# qa fields worth surfacing in the Eval tab (heavy retriever artifacts dropped).
QA_KEEP = [
    "question", "answer", "exe_ans", "program", "program_re", "steps",
    "gold_inds", "explanation", "ann_table_rows", "ann_text_rows",
]


def safe_key(filename: str) -> str:
    """`ADI/2009/page_49.pdf` -> `ADI__2009__page_49.pdf` (filesystem-safe)."""
    return filename.replace("/", "__")


def main() -> None:
    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    default_src = os.path.join(
        os.environ.get("DATA_ROOT", "/mnt/ramdisk/blobstore/timchen0618/data"),
        "finqa", "dataset",
    )
    if not os.path.isdir(default_src):
        default_src = os.path.join(root, "dataset")

    ap = argparse.ArgumentParser()
    ap.add_argument("--src", default=default_src, help="dir with the raw split json files")
    args = ap.parse_args()

    records_dir = os.path.join(root, "corpus", "records")
    os.makedirs(records_dir, exist_ok=True)
    # clear stale shards
    for fn in os.listdir(records_dir):
        if fn.endswith(".json"):
            os.remove(os.path.join(records_dir, fn))

    corpus: dict[str, dict] = {}
    counts: dict[str, Counter] = defaultdict(Counter)
    eval_rows: list[dict] = []

    for split in SPLITS:
        path = os.path.join(args.src, f"{split}.json")
        with open(path) as f:
            data = json.load(f)
        for x in data:
            fn = x["filename"]
            counts[fn][split] += 1
            if fn not in corpus:
                ticker, year, page = fn.split("/")
                corpus[fn] = {
                    "filename": fn,
                    "ticker": ticker,
                    "year": year,
                    "page": page.replace("page_", "").replace(".pdf", ""),
                    "pre_text": x["pre_text"],
                    "table": x["table"],
                    "post_text": x["post_text"],
                }
            qa = x.get("qa", {})
            row = {"id": x["id"], "filename": fn, "split": split}
            for k in QA_KEEP:
                if k in qa:
                    row[k] = qa[k]
            eval_rows.append(row)

    # eval.jsonl
    with open(os.path.join(root, "eval.jsonl"), "w") as f:
        for row in eval_rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")

    # corpus records + index
    index = []
    for fn, doc in corpus.items():
        doc["num_questions"] = sum(counts[fn].values())
        doc["splits"] = sorted(counts[fn].keys())
        key = safe_key(fn)
        with open(os.path.join(records_dir, key + ".json"), "w") as f:
            json.dump(doc, f, ensure_ascii=False)
        index.append({
            "filename": fn,
            "key": key,
            "ticker": doc["ticker"],
            "year": doc["year"],
            "page": doc["page"],
            "num_questions": doc["num_questions"],
            "splits": doc["splits"],
        })

    def sort_key(d):
        return (d["ticker"], d["year"], int(d["page"]) if d["page"].isdigit() else 0)

    index.sort(key=sort_key)
    with open(os.path.join(root, "corpus", "index.json"), "w") as f:
        json.dump(index, f, ensure_ascii=False)

    print(f"eval.jsonl: {len(eval_rows)} rows")
    print(f"corpus: {len(index)} docs -> corpus/records/*.json + corpus/index.json")


if __name__ == "__main__":
    main()