#!/usr/bin/env python3
"""Build first letters' model files and their catalogue.

First letters (abbreviation expansion) turns the first letter of each word
the person types ("iwtgh") into whole sentences ("I want to go home"), on the
device: Qwen3-0.6B-Base, fine-tuned for it (LoRA, merged), stored in 4 bits
(INT4, block 16) and run on onnxruntime-web's WebGPU execution provider. The
model was made and tested in the model library (README.md says where, and
how). This script takes its browser bundle (onnx/ae-en/), checks each file
against the size and SHA-256 pinned here (FILES), names each by its contents,
splits the one file larger than 64 MiB into parts (it holds one tensor, the
table of word embeddings, so the bundle could not split it), and writes the
catalogue. Every part is at most 64 MiB, so the app's model store can check
each one in memory on a tablet and Cloudflare can cache it.

Outputs (schema "aac-ae-index/1", read by crates/aac-web/src/abbrev.rs):

  ../aac-ae-packs/                      beside the open-standards checkout ($AAC_AE_DIR
                                        overrides; --out too): what the site builds copy
                                        to /app/ae/ after tools/ae/check_packs.py
    index.json                          pointer {schema, current}; serve no-store
    index.<sha12>.json                  the immutable catalogue
    LICENSES.md                         the licences and attribution
    licences/<name>.<sha12>.txt         the licence texts
    <model-id>/<file>.<sha12>.<ext>     the graph, its weights in parts, the tokenizer,
                                        the model's configuration and letter tables
  data/ae/index.json                    committed copy of the catalogue
  data/ae/LICENSES.md                   committed copy of LICENSES.md

The languages, their status and the evidence for each come from
tools/ae/languages.json. A language is ready only when we measured it on this
app's own worker (tools/ae/eval) with the sentence meant among the five
offered for at least MIN_TOP5 percent of at least MIN_ITEMS test replies.

Usage (from aac/aac-rs):
  python3 tools/ae/build_ae.py --source DIR     # DIR: the model library's onnx/ae-en
  python3 tools/ae/build_ae.py                  # the same, from $AAC_AE_SOURCE or tools/ae/.cache/ae-en
  python3 tools/ae/build_ae.py --out DIR --no-commit
                                                # a trial build elsewhere, data/ae left alone

Standard library only (Python 3.9 or later).
"""

import argparse
import hashlib
import json
import math
import os
import re
import sys
from datetime import datetime, timezone
from pathlib import Path

HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[1]  # aac/aac-rs
SOURCE_DEFAULT = Path(os.environ.get("AAC_AE_SOURCE") or HERE / ".cache" / "ae-en")
OUT_DEFAULT = Path(os.environ.get("AAC_AE_DIR") or ROOT.parents[2] / "aac-ae-packs")
COMMITTED = ROOT / "data" / "ae"
SCHEMA = "aac-ae-index/1"
# The runtime the graph is for: crates/aac-web/src/abbrev.rs RUNTIME_ID (the
# listener's onnxruntime-web, shared).
RUNTIME = "onnxruntime-web@1.30.0"
PART_BYTES = 64 * 1024 * 1024
# A language is ready only when the sentence meant was among the five offered
# for at least this percent of at least MIN_ITEMS test replies, measured on
# this app's own worker: abbrev.rs MIN_TOP5 and MIN_ITEMS.
MIN_TOP5 = 50.0
MIN_ITEMS = 200
# The day the sources were read: the catalogue's date, and every file's
# modification time, so a rebuild of the same sources is byte- and
# time-identical (rsync then sends nothing).
RETRIEVED = "2026-10-01"
EPOCH = int(datetime.fromisoformat(RETRIEVED).replace(tzinfo=timezone.utc).timestamp())

MODEL_ID = "ae-en-qwen3-0.6b"
MODEL = {
    "name": "Qwen3 0.6B, taught first letters (English)",
    "source": ("https://huggingface.co/Qwen/Qwen3-0.6B-Base/tree/da87bfb608c14b7cf20ba1ce41287e8de496c0cd, "
               "fine-tuned by AuraSense to turn first letters into sentences (LoRA rank 16 on every linear layer, "
               "merged; 281,414 examples), exported to INT4 (block 16) for WebGPU with onnxruntime-genai 0.17.1, "
               "its graph edited so every node runs on WebGPU"),
    # GPU memory while loaded: the weights (364 MiB) and the search's caches
    # (about 28 MB at beam 8), as the model library's LOG.md counts them.
    "gpu_memory_bytes": 420_000_000,
    "min_memory_gb": 4,
}

# Every file of the bundle the app runs, as the model library's export names
# it (onnx/ae-en/ae_manifest.json), by size and SHA-256.
FILES = {
    "model.onnx": (260299, "70208fbe674ebb8b020c6cdb0692d8a4111e513a516c3d253030b1b4137b29cd"),
    "model.onnx.data.00": (67094528, "db4f784defdbf17f7d54da01b610a433cf1309b66aa0c655b4140ecd8dece071"),
    "model.onnx.data.01": (66846720, "301e5ce3188091724466ab5b932c0eb0b4b49fab4356b422c96001c601442b8c"),
    "model.onnx.data.02": (66846720, "16f8ec3c0935fb95fa6e62b10996ab822b19c09387f75559792195f1447c3b0c"),
    "model.onnx.data.03": (66977792, "4523f62019769e8e5d35e923ca2029c84750a716654af0f6082c3f695fce574e"),
    "model.onnx.data.04": (15990784, "511489abbea546330a8f41af2f4754dfd935c436bda41e18f6be19fc64996255"),
    "model.onnx.data.05": (77791232, "4a5d25081ed0236ff3a08bc3c7d71ac1e03a2e0f527e81b2aae101fceabfb786"),
    "model.onnx.data.06": (19447808, "86d3ec9fe2c8f9ede6657d5f1c1574798bb1afb198b1dff300f41a028b2cef2d"),
    "tokenizer.json": (11422650, "be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506"),
    "tokenizer_config.json": (731, "6bd46c5753c8cfda8c8b5b3df8c6351643ca5c61049eba293ccd05a004ea441f"),
    "genai_config.json": (1836, "628a1efa3ca8856a86c765afbf8c6be5c6337ce13de3adad20d292fbec3af36b"),
    "ae_vocab.json": (30080, "14a4b5bbd3fe5b1cbf0be974f38b0f0736bdc659d6ca9fb7c7a06c1dcd464e9b"),
    "ae_vocab.bin": (759680, "4240f3f07ab5c9d9b599f24d2fb2044387b9daace4022c94b9e6d9eec5fd6d9f"),
}
GRAPH = "model.onnx"
# The graph's weights files, in the order the graph names them.
DATA = [f"model.onnx.data.{i:02d}" for i in range(7)]
# The other files, by the field the catalogue lists each under.
OTHERS = {
    "tokenizer": "tokenizer.json",
    "tokenizer_config": "tokenizer_config.json",
    "config": "genai_config.json",
    "vocab": "ae_vocab.json",
    "vocab_table": "ae_vocab.bin",
}

# The licence texts, verbatim, by SHA-256: Qwen3-0.6B-Base's LICENSE (Apache
# License 2.0, Alibaba Cloud), as in its repository at the revision above.
LICENCES = {
    "Apache-2.0-Qwen3": (HERE / "licences" / "Apache-2.0-Qwen3.txt",
                         "832dd9e00a68dd83b3c3fb9f5588dad7dcf337a0db50f7d9483f310cd292e92e"),
}
CC_BY_4 = "https://creativecommons.org/licenses/by/4.0/"
CREDITS = [
    {
        "what": "the model",
        "attribution": ("Qwen3-0.6B-Base by the Qwen team (Alibaba Cloud), fine-tuned by AuraSense to turn first letters "
                        "into sentences and stored in 4 bits for WebGPU"),
        "url": "https://huggingface.co/Qwen/Qwen3-0.6B-Base",
        "licence": "Apache-2.0",
        "licence_url": "https://www.apache.org/licenses/LICENSE-2.0",
        "licence_key": "Apache-2.0-Qwen3",
    },
    {
        "what": "what it learned from",
        "attribution": ("Turk AAC communications by Keith Vertanen and Per Ola Kristensson (“The Imagination of "
                        "Crowds: Conversational AAC Language Modeling using Crowdsourcing and Large Data Sources”, "
                        "EMNLP 2011), turned into first-letter examples"),
        "url": "https://aactext.org/imagine/",
        "licence": "CC BY 4.0",
        "licence_url": CC_BY_4,
    },
    {
        "what": "what it learned from",
        "attribution": ("Synthetic-Persona-Chat by Pegah Jandaghi, XiangHai Sheng, Xinyi Bai, Jay Pujara and Hakim "
                        "Sidahmed (Google, 2023), its consecutive turns turned into first-letter examples"),
        "url": "https://huggingface.co/datasets/google/Synthetic-Persona-Chat",
        "licence": "CC BY 4.0",
        "licence_url": CC_BY_4,
    },
    {
        "what": "what it learned from",
        "attribution": ("Conversations written for it by AuraSense with Gemma 4 26B-A4B (Google), whose outputs are "
                        "free to use"),
        "url": "https://ai.google.dev/gemma",
        "licence": "Apache-2.0",
        "licence_url": "https://www.apache.org/licenses/LICENSE-2.0",
    },
]


class BuildError(Exception):
    pass


def log(msg):
    print(msg, flush=True)


def sha256_bytes(data):
    return hashlib.sha256(data).hexdigest()


def sha256_file(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 22), b""):
            h.update(chunk)
    return h.hexdigest()


def dumps(obj):
    return (json.dumps(obj, indent=1, ensure_ascii=False) + "\n").encode("utf-8")


def languages():
    doc = json.loads((HERE / "languages.json").read_text(encoding="utf-8"))
    langs = doc["languages"]
    for code, l in langs.items():
        if not re.fullmatch(r"[a-z]{2,3}", code):
            raise BuildError(f"languages.json: code {code!r}")
        if not l.get("name") or not l.get("note"):
            raise BuildError(f"languages.json: {code} needs a name and a note")
        if l["status"] == "ready":
            if l.get("model") != MODEL_ID:
                raise BuildError(f"languages.json: {code} is ready with an unknown model {l.get('model')!r}")
            if l.get("measured") is None or l["measured"] < MIN_TOP5 or (l.get("items") or 0) < MIN_ITEMS:
                raise BuildError(f"languages.json: {code} is ready without our measured top 5 of at least {MIN_TOP5} "
                                 f"on at least {MIN_ITEMS} replies")
            if not l.get("source"):
                raise BuildError(f"languages.json: {code} is ready without saying where its numbers come from")
        elif l["status"] != "not-yet":
            raise BuildError(f"languages.json: {code} status {l['status']!r}")
    return dict(sorted(langs.items()))


def check_sources(source):
    """Every pinned file in `source`, with its size and SHA-256 checked."""
    found = {}
    for name, (size, digest) in FILES.items():
        path = source / name
        if not path.is_file():
            raise BuildError(f"{path} is missing")
        if path.stat().st_size != size:
            raise BuildError(f"{path} is {path.stat().st_size} bytes; the pin is {size}")
        got = sha256_file(path)
        if got != digest:
            raise BuildError(f"{path} has sha256 {got}; the pin is {digest}")
        found[name] = path
    # The graph names its weights files by these names (ONNX external data
    # locations are plain strings in the graph).
    blob = found[GRAPH].read_bytes()
    for d in DATA:
        if d.encode() not in blob:
            raise BuildError(f"{GRAPH} does not name {d}")
    return found


def split_points(size):
    """Where a file is cut into parts of at most PART_BYTES: as few parts as
    that allows, of near-equal size (multiples of 4 KiB but the last)."""
    n = max(1, math.ceil(size / PART_BYTES))
    step = math.ceil(size / n / 4096) * 4096
    return [(i * step, min(size, (i + 1) * step)) for i in range(n)]


def plan(found):
    """Every file to write, as rel -> ("copy", src) | ("slice", (src, start, end)) | ("data", bytes),
    and the catalogue."""
    files = {}
    licence_rel = {}
    for key, (path, digest) in LICENCES.items():
        data = path.read_bytes()
        if sha256_bytes(data) != digest:
            raise BuildError(f"{path} is not the licence text pinned here")
        rel = f"licences/{key}.{digest[:12]}.txt"
        files[rel] = ("data", data)
        licence_rel[key] = rel

    def named(stem, digest, ext):
        return f"{MODEL_ID}/{stem}.{digest[:12]}.{ext}"

    def whole(name, stem, ext):
        size, digest = FILES[name]
        rel = named(stem, digest, ext)
        files[rel] = ("copy", found[name])
        return {"url": rel, "bytes": size, "sha256": digest}

    data = []
    for i, name in enumerate(DATA):
        size, digest = FILES[name]
        cuts = split_points(size)
        if len(cuts) == 1:
            parts = [whole(name, f"data-{i:02d}", "bin")]
        else:
            parts = []
            with open(found[name], "rb") as f:
                for k, (start, end) in enumerate(cuts, 1):
                    f.seek(start)
                    chunk = f.read(end - start)
                    d = sha256_bytes(chunk)
                    rel = named(f"data-{i:02d}-{k}", d, "bin")
                    files[rel] = ("slice", (found[name], start, end))
                    parts.append({"url": rel, "bytes": end - start, "sha256": d})
        if sum(p["bytes"] for p in parts) != size:
            raise BuildError(f"{name}: its parts do not add up")
        data.append({"name": name, "bytes": size, "parts": parts})

    model = {
        "name": MODEL["name"],
        "runtime": RUNTIME,
        "source": MODEL["source"],
        "bytes": sum(s for s, _ in FILES.values()),
        "graph": whole(GRAPH, "model", "onnx"),
        "data": data,
    }
    for field, name in OTHERS.items():
        stem, _, ext = name.rpartition(".")
        model[field] = whole(name, stem, ext)
    model["credits"] = [
        {k: v for k, v in c.items() if k != "licence_key"}
        | ({"licence_file": licence_rel[c["licence_key"]]} if "licence_key" in c else {})
        for c in CREDITS
    ]
    model["gpu_memory_bytes"] = MODEL["gpu_memory_bytes"]
    model["min_memory_gb"] = MODEL["min_memory_gb"]
    catalogue = {"schema": SCHEMA, "generated": RETRIEVED, "models": {MODEL_ID: model}, "languages": languages()}
    cat_bytes = dumps(catalogue)
    cat_name = f"index.{sha256_bytes(cat_bytes)[:12]}.json"
    files[cat_name] = ("data", cat_bytes)
    files["index.json"] = ("data", dumps({"schema": SCHEMA, "current": cat_name}))
    files["LICENSES.md"] = ("data", licences_md(catalogue).encode("utf-8"))
    return files, cat_bytes


def licences_md(catalogue):
    m = catalogue["models"][MODEL_ID]
    ready = [f"{l['name']} ({c})" for c, l in catalogue["languages"].items() if l["status"] == "ready"]
    n = sum(len(d["parts"]) for d in m["data"]) + 1 + len(OTHERS)
    out = [
        "# First letters: licences\n\n"
        "First letters in AuraSense AAC (aac-web) downloads one small language model, on request, and runs it on "
        "the person's device, on its graphics chip: it turns the first letter of each word they type into whole "
        "sentences for them to choose from. Nothing typed is sent anywhere.\n",
        f"\n## {m['name']}\n\n",
    ]
    for c in m["credits"]:
        text = f"; the text is `{c['licence_file']}`" if c.get("licence_file") else ""
        page = f" ({c['url']})" if c.get("url") else ""
        out.append(f"- {c['what'].capitalize()}: {c['attribution']}{page}. Licence: {c['licence']} ({c['licence_url']}){text}.\n")
    out.append(
        f"- Source: {m['source']}.\n"
        "- What was changed: the model's weights by the fine-tuning above; the training data's sentences were "
        "turned into examples (each reply's first letters, what the partner said before it, some with typing "
        "slips), and nothing of them is shipped. The graph was edited so every node runs on WebGPU in "
        "onnxruntime-web 1.30 (its attention-mask arithmetic fed in directly, the search's caches reordered on the "
        f"graphics chip). {n} files of at most 64 MiB, each named by its own SHA-256 ({m['bytes'] / 1e6:.0f} MB).\n"
        "- The way the letters are written (each word's first letter, contractions split at the apostrophe, "
        "punctuation where it was typed) follows S. Cai et al., “Context-Aware Abbreviation Expansion Using Large "
        "Language Models”, NAACL 2022 (the method only; none of its data is used).\n"
        "- Tested on Turk Dialogues (K. Vertanen, 2017), which is used for testing only and is not shipped.\n"
        f"- Ready in: {', '.join(ready) or 'no language yet'}.\n"
        "\nThe runtime that runs it (onnxruntime-web, MIT) is the listener's, vendored with its licence, notice and "
        "ONNX Runtime's third-party notices under `vendor/onnxruntime-web-<version>-<revision>/`; the tokenizer "
        "library (@huggingface/tokenizers, Apache-2.0) and the app's worker are vendored with their licence and "
        "notice under `vendor/hf-tokenizers-<version>-<revision>/`.\n"
    )
    return "".join(out)


def write_tree(out, files):
    if out.exists() and any(out.iterdir()) and not (out / "index.json").exists():
        raise BuildError(f"refusing to prune non-empty {out}: it has no index.json from an earlier build")
    out.mkdir(parents=True, exist_ok=True)

    # Files and licences first, the catalogue after them, the pointer last: a
    # run that stops part way never leaves a pointer to files not there yet.
    def order(rel):
        if rel == "index.json":
            return (3, rel)
        if rel == "LICENSES.md":
            return (2, rel)
        return (1 if re.fullmatch(r"index\.[0-9a-f]{12}\.json", rel) else 0, rel)

    def put(dest, chunks):
        tmp = dest.with_name(dest.name + ".tmp")
        with open(tmp, "wb") as o:
            for chunk in chunks:
                o.write(chunk)
        os.replace(tmp, dest)

    def read(src, start=0, end=None):
        with open(src, "rb") as f:
            f.seek(start)
            left = (end - start) if end is not None else None
            while left is None or left > 0:
                chunk = f.read(1 << 22 if left is None else min(1 << 22, left))
                if not chunk:
                    break
                if left is not None:
                    left -= len(chunk)
                yield chunk

    for rel in sorted(files, key=order):
        kind, spec = files[rel]
        dest = out / rel
        dest.parent.mkdir(parents=True, exist_ok=True)
        if kind == "copy":
            size = spec.stat().st_size
            if not (dest.exists() and dest.stat().st_size == size and sha256_file(dest) == sha256_file(spec)):
                put(dest, read(spec))
        elif kind == "slice":
            src, start, end = spec
            want = rel.rsplit(".", 2)[-2]  # the name carries its own hash
            if not (dest.exists() and dest.stat().st_size == end - start and sha256_file(dest).startswith(want)):
                put(dest, read(src, start, end))
        elif not (dest.exists() and dest.read_bytes() == spec):
            put(dest, [spec])
        os.chmod(dest, 0o644)
        os.utime(dest, (EPOCH, EPOCH))
    keep = set(files)
    for p in sorted(out.rglob("*"), reverse=True):
        rel = p.relative_to(out).as_posix()
        if p.is_file() and rel not in keep:
            p.unlink()
            log(f"prune: {rel}")
        elif p.is_dir() and not any(p.iterdir()):
            p.rmdir()


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--source", type=Path, default=SOURCE_DEFAULT,
                    help="the model library's browser bundle (onnx/ae-en)")
    ap.add_argument("--out", type=Path, default=OUT_DEFAULT)
    ap.add_argument("--no-commit", action="store_true", help="leave data/ae alone (a trial build)")
    args = ap.parse_args()
    try:
        if args.no_commit and args.out == OUT_DEFAULT:
            raise BuildError("a trial build (--no-commit) goes elsewhere: give --out DIR")
        found = check_sources(args.source)
        files, cat_bytes = plan(found)
        write_tree(args.out, files)
        if not args.no_commit:
            COMMITTED.mkdir(parents=True, exist_ok=True)
            (COMMITTED / "index.json").write_bytes(cat_bytes)
            (COMMITTED / "LICENSES.md").write_bytes(files["LICENSES.md"][1])
    except BuildError as e:
        print(f"✗ {e}", file=sys.stderr)
        return 1
    n = sum(1 for kind, _ in files.values() if kind in ("copy", "slice"))
    log(f"✓ {args.out}: {MODEL_ID}, {n} files, catalogue {json.loads(files['index.json'][1])['current']}")
    return 0


if __name__ == "__main__":
    sys.exit(main())