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

Reads the two source QA files (simple.json = concise final answers,
complex.json = longer analytical final answers) and derives:

- eval.json          one entry per question (dataset-tagged); question, gold
                     final answer, supporting facts (source_answer) and the
                     supporting-document list (each with the value_* fields that
                     question used + their schema descriptions).
- corpus_index.json  one entry per unique document (589 == the PDF corpus).
                     Title = "symbol · year · doctype"; value_* fields are the
                     union observed across every question that cites the doc;
                     `pdf` points at the HF dataset CDN (streamed, not bundled).

Run from the viewer repo root:
    python scripts/build_data.py

Reads:  simple.json, complex.json
Writes: eval.json, corpus_index.json
"""
import json
import os
from collections import OrderedDict

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SOURCES = [("simple", os.path.join(ROOT, "simple.json")),
           ("complex", os.path.join(ROOT, "complex.json"))]
EVAL_OUT = os.path.join(ROOT, "eval.json")
CORPUS_OUT = os.path.join(ROOT, "corpus_index.json")

# PDFs live on the Hugging Face dataset; the viewer streams them from the CDN
# so the Space never has to bundle the ~4 GB corpus.
PDF_URL = "https://huggingface.co/datasets/Zhanli-Li/MuDABench/resolve/main/data/pdf/{id}.pdf"


def value_fields(meta):
    """Return [(field_name, value), ...] for the value_* keys in a metadata row."""
    out = []
    for k, v in meta.items():
        if k.startswith("value_"):
            out.append((k, v))
    return out


def doc_title(symbol, year, doctype):
    parts = [str(p) for p in (symbol, year, doctype) if p not in (None, "")]
    return " · ".join(parts) if parts else "(untitled)"


def main():
    eval_rows = []
    # corpus[id] accumulates the merged view of a document across all questions.
    corpus = OrderedDict()

    for dataset, path in SOURCES:
        with open(path, encoding="utf-8") as f:
            data = json.load(f)
        for q in data:
            qid = q.get("question_id")
            src = q.get("source_answer")
            if isinstance(src, str):
                src = [src] if src.strip() else []
            elif not isinstance(src, list):
                src = []

            docs = []
            for meta in q.get("metadata", []):
                did = meta.get("id")
                symbol = meta.get("symbol")
                year = meta.get("year")
                doctype = meta.get("doctype")
                schema = meta.get("schema", {}) or {}
                vfields = value_fields(meta)

                docs.append({
                    "id": did,
                    "title": doc_title(symbol, year, doctype),
                    "symbol": symbol,
                    "year": year,
                    "doctype": doctype,
                    "fields": [
                        {"name": n, "value": v, "desc": schema.get(n, "")}
                        for n, v in vfields
                    ],
                })

                if did is None:
                    continue
                c = corpus.get(did)
                if c is None:
                    c = corpus[did] = {
                        "id": did,
                        "symbol": symbol,
                        "year": year,
                        "doctype": doctype,
                        "title": doc_title(symbol, year, doctype),
                        "_fields": OrderedDict(),   # name -> {"desc", "values"[]}
                        "referenced_by": [],
                    }
                for n, v in vfields:
                    slot = c["_fields"].get(n)
                    if slot is None:
                        slot = c["_fields"][n] = {"desc": schema.get(n, ""), "values": []}
                    if not slot["desc"] and schema.get(n):
                        slot["desc"] = schema[n]
                    if v not in slot["values"]:
                        slot["values"].append(v)
                c["referenced_by"].append({"qid": qid, "dataset": dataset})

            eval_rows.append({
                "qid": qid,
                "dataset": dataset,
                "question": q.get("question", ""),
                "final_answer": q.get("final_answer", ""),
                "source_answer": src,
                "docs": docs,
            })

    # finalize corpus entries
    corpus_rows = []
    for c in corpus.values():
        fields = [
            {"name": n, "desc": slot["desc"], "values": slot["values"]}
            for n, slot in c["_fields"].items()
        ]
        corpus_rows.append({
            "id": c["id"],
            "title": c["title"],
            "symbol": c["symbol"],
            "year": c["year"],
            "doctype": c["doctype"],
            "fields": fields,
            "n_refs": len(c["referenced_by"]),
            "referenced_by": c["referenced_by"],
            "pdf": PDF_URL.format(id=c["id"]),
        })

    corpus_rows.sort(key=lambda d: (
        str(d["doctype"] or ""), str(d["symbol"] or ""), str(d["year"] or "")
    ))

    # collision report (docs sharing an identical title)
    seen = {}
    collisions = 0
    for d in corpus_rows:
        seen.setdefault(d["title"], []).append(d["id"])
    for title, ids in seen.items():
        if len(ids) > 1:
            collisions += 1

    with open(EVAL_OUT, "w", encoding="utf-8") as f:
        json.dump(eval_rows, f, ensure_ascii=False, indent=0)
    with open(CORPUS_OUT, "w", encoding="utf-8") as f:
        json.dump(corpus_rows, f, ensure_ascii=False, indent=0)

    print(f"wrote {EVAL_OUT}: {len(eval_rows)} questions")
    print(f"wrote {CORPUS_OUT}: {len(corpus_rows)} documents "
          f"({collisions} title-collision groups)")
    from collections import Counter
    dt = Counter(d["doctype"] for d in corpus_rows)
    print("doctypes:", dict(dt))


if __name__ == "__main__":
    main()