#!/usr/bin/env python3
"""Write corpus.tar.xz and corpus.json: the text the dictionary measurement reads
(research/text-dictionary.md §2), copied from this repository's own files under wai/, which
wai/LICENSE covers.

    python3 make_corpus.py          # from anywhere inside a checkout

Four families, each object one payload a sink could fetch:

- `notes`: the changelog's entries, one object each (the files of wai/changes/ and each `### `
  section of wai/CHANGELOG.md), newest first.
- `spec`: wai/extensions/*.md, one object a document.
- `code`: wai/wai-rs/src/**/*.rs of at most 32 KiB.
- `records`: the JSON files of at most 8 KiB in the conformance corpora named in RECORD_DIRS,
  whose group is the directory they sit in.

An object whose bytes repeat an earlier object of its family is dropped. The split is fixed
here, before anything is measured, and no held-out object or group trains a dictionary:
`notes` by time, the newest quarter held out, as a dictionary meets what is written after it;
`spec` and `code` by the digest of the path; `records` by the digest of the group, so that
every near-duplicate of a held-out record is held out with it.

The archive is byte for byte reproducible: entries sorted, no times, no owners, xz preset 9e.
"""
import hashlib
import io
import json
import lzma
import pathlib
import re
import subprocess
import tarfile

HERE = pathlib.Path(__file__).resolve().parent
ROOT = pathlib.Path(subprocess.check_output(["git", "-C", str(HERE), "rev-parse", "--show-toplevel"], text=True).strip())
RECORD_DIRS = ["vector-conformance", "pack-conformance", "http-conformance", "live-conformance",
               "encode-conformance", "staged-conformance", "telemetry-conformance"]
MAX_RECORD, MAX_CODE = 8192, 32768


def sha(b):
    return hashlib.sha256(b).hexdigest()


def tracked():
    out = subprocess.check_output(["git", "-C", str(ROOT), "ls-files", "-z", "wai"])
    return sorted(p for p in out.decode().split("\0") if p)


def quarter(key):
    return int(sha(key.encode())[:8], 16) % 4 == 0


def notes():
    objs = []
    frag = sorted((p for p in tracked() if re.fullmatch(r"wai/changes/\d{4}-\d\d-\d\d-.*\.md", p)), reverse=True)
    for p in frag:
        objs.append((p, (ROOT / p).read_bytes()))
    text = (ROOT / "wai/CHANGELOG.md").read_bytes()
    parts = re.split(rb"(?m)^(?=### )", text)
    for i, part in enumerate(parts[1:]):
        # A release heading after an entry belongs to no entry.
        part = re.split(rb"(?m)^## ", part)[0]
        objs.append((f"wai/CHANGELOG.md#{i}", part))
    cut = len(objs) // 4
    return [(s, b, None, i < cut) for i, (s, b) in enumerate(objs)]


def by_path(paths):
    return [(p, (ROOT / p).read_bytes(), None, quarter(p)) for p in paths]


def spec():
    return by_path([p for p in tracked() if re.fullmatch(r"wai/extensions/[^/]*\.md", p)])


def code():
    src = [p for p in tracked() if p.startswith("wai/wai-rs/src/") and p.endswith(".rs")]
    return by_path([p for p in src if (ROOT / p).stat().st_size <= MAX_CODE])


def records():
    out = []
    for p in tracked():
        top = p.split("/")[1] if p.count("/") >= 2 else ""
        if top in RECORD_DIRS and p.endswith(".json") and 0 < (ROOT / p).stat().st_size <= MAX_RECORD:
            group = p.rsplit("/", 1)[0]
            out.append((p, (ROOT / p).read_bytes(), group, quarter(group)))
    return out


def main():
    manifest, members = [], []
    for family, objs in (("notes", notes()), ("spec", spec()), ("code", code()), ("records", records())):
        seen = set()
        n = 0
        for source, body, group, held in objs:
            digest = sha(body)
            if digest in seen or not body:
                continue
            seen.add(digest)
            name = f"{family}/{n:05d}"
            n += 1
            entry = {"name": name, "family": family, "source": source, "bytes": len(body), "sha256": digest,
                     "split": "held-out" if held else "train"}
            if group:
                entry["group"] = group
            manifest.append(entry)
            members.append((name, body))
    buf = io.BytesIO()
    with tarfile.open(fileobj=buf, mode="w", format=tarfile.USTAR_FORMAT) as tar:
        for name, body in members:
            ti = tarfile.TarInfo(name)
            ti.size = len(body)
            tar.addfile(ti, io.BytesIO(body))
    packed = lzma.compress(buf.getvalue(), preset=9 | lzma.PRESET_EXTREME)
    (HERE / "corpus.tar.xz").write_bytes(packed)
    head = subprocess.check_output(["git", "-C", str(ROOT), "rev-parse", "HEAD"], text=True).strip()
    doc = {"format": "wai.research.text-dictionary.corpus/1", "from_commit": head, "licence": "Apache-2.0 (wai/LICENSE)",
           "archive_sha256": sha(packed), "objects": manifest}
    (HERE / "corpus.json").write_text(json.dumps(doc, indent=0) + "\n")
    print(f"{len(manifest)} objects, {sum(e['bytes'] for e in manifest)} bytes, archive {len(packed)} bytes")
    for family in ("notes", "spec", "code", "records"):
        es = [e for e in manifest if e["family"] == family]
        held = [e for e in es if e["split"] == "held-out"]
        print(f"  {family}: {len(es)} objects, {sum(e['bytes'] for e in es)} bytes; held out {len(held)}, {sum(e['bytes'] for e in held)} bytes")


if __name__ == "__main__":
    main()