#!/usr/bin/env python3
"""Build the sentence helper's model parts and their catalogue.

The models are Google's web builds of Gemma 4 for LiteRT-LM (E2B, and E4B
for languages E2B is not good enough in), each pinned by Hugging Face
revision, size and SHA-256 (MODELS below). Each is split into parts of at most
64 MiB, named by their contents, so the app's model store can check each part
in memory on a tablet and Cloudflare can cache every part. A catalogue holds
the models some ready language uses (or those named with --models, for an
evaluation build).

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

  ../aac-ai-packs/                      beside the open-standards checkout ($AAC_AI_DIR
                                        overrides; --out too): what the site builds copy
                                        to /app/ai/ after tools/ai/check_packs.py
    index.json                          pointer {schema, current}; serve no-store
    index.<sha12>.json                  the immutable catalogue
    LICENSES.md                         the model's licence and attribution
    licences/Apache-2.0.<sha12>.txt     the licence text
    <model-id>/part-NN.<sha12>.bin
  data/ai/index.json                    committed copy of the catalogue
  data/ai/LICENSES.md                   committed copy of LICENSES.md

The languages, their status and the examples the model is shown come from
tools/ai/languages.json. A language is listed there only after testing.

Usage (from aac/aac-rs):
  python3 tools/ai/build_ai.py                           # download pinned models once, build
  python3 tools/ai/build_ai.py --source ID=FILE ...      # use local copies (their pins are checked)
  python3 tools/ai/build_ai.py --offline                 # cached downloads only
  python3 tools/ai/build_ai.py --models ID --out DIR --no-commit
                                                         # an evaluation build of other models

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

import argparse
import hashlib
import json
import os
import re
import shutil
import sys
import urllib.request
from datetime import datetime, timezone
from pathlib import Path

HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[1]  # aac/aac-rs
CACHE = HERE / ".cache"
OUT_DEFAULT = Path(os.environ.get("AAC_AI_DIR") or ROOT.parents[2] / "aac-ai-packs")
COMMITTED = ROOT / "data" / "ai"
SCHEMA = "aac-ai-index/1"
# The runtime the parts are for: crates/aac-web/src/ai_helper.rs RUNTIME_ID.
RUNTIME = "litert-lm@0.17.1"
PART_BYTES = 64 * 1024 * 1024
# 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())

ATTRIBUTION = ("{name} by Google DeepMind, in Google AI Edge's web build for LiteRT-LM "
               "(litert-community on Hugging Face)")
MODELS = {
    "gemma-4-e2b-it-web": {
        "repo": "litert-community/gemma-4-E2B-it-litert-lm",
        "revision": "b3ca0d2f076785a8f4b2219ddbd2bdb99954eae1",
        "file": "gemma-4-E2B-it-web.litertlm",
        "bytes": 2008432640,
        "sha256": "3a08e8d94e23b814ae5414469c370c503813949acb8ceaa17e4ebf8a35af35b5",
        "name": "Gemma 4 E2B",
        # What it held in GPU memory when measured (Chrome on a Mac, an iPad
        # Air, 2026-09-30): about 1.8 GB.
        "gpu_memory_bytes": 1_800_000_000,
        "min_memory_gb": 4,
        "covers": [],
    },
    "gemma-4-e4b-it-web": {
        "repo": "litert-community/gemma-4-E4B-it-litert-lm",
        "revision": "2eee7ac325f20eb8c9ac1d0e972f7c84663062da",
        "file": "gemma-4-E4B-it-web.litertlm",
        "bytes": 2969059328,
        "sha256": "3904d826d5dddd25ea173e85204caec09e68ba038116e9b992b69cbdc94f57a0",
        "name": "Gemma 4 E4B",
        # Measured in Chrome on a Mac, 2026-10-01: a 3.0 GB GPU process.
        "gpu_memory_bytes": 3_000_000_000,
        "min_memory_gb": 8,
        # The larger model serves every language the smaller one is ready in.
        "covers": ["gemma-4-e2b-it-web"],
    },
}
LICENCE = "Apache-2.0"
LICENCE_URL = "https://ai.google.dev/gemma/docs/gemma_4_license"
CONTEXT_TOKENS = 1024
# The Apache License 2.0 text, as LiteRT-LM ships it (third_party/litert-lm/LICENSE.txt).
APACHE = ROOT / "third_party" / "litert-lm" / "LICENSE.txt"
APACHE_SHA256 = "c71d239df91726fc519c6eb72d318ec65820627232b2f796219e87dcf35d0ab4"


class BuildError(Exception):
    pass


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


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 model_source(model, source, offline):
    """A pinned model's file: its --source, else the cached download, else
    download it (into the cache, checked before it is used)."""
    def check(path):
        size = path.stat().st_size
        if size != model["bytes"]:
            raise BuildError(f"{path} is {size} bytes; the pin is {model['bytes']}")
        got = sha256_file(path)
        if got != model["sha256"]:
            raise BuildError(f"{path} has sha256 {got}; the pin is {model['sha256']}")
        return path

    if source:
        return check(Path(source))
    cached = CACHE / model["revision"] / model["file"]
    if cached.exists():
        return check(cached)
    if offline:
        raise BuildError(f"{cached} is not cached and --offline was given")
    url = f"https://huggingface.co/{model['repo']}/resolve/{model['revision']}/{model['file']}"
    log(f"download: {url}")
    cached.parent.mkdir(parents=True, exist_ok=True)
    tmp = cached.with_name(cached.name + ".tmp")
    with urllib.request.urlopen(url) as res, open(tmp, "wb") as out:
        shutil.copyfileobj(res, out, 1 << 22)
    check(tmp)
    os.replace(tmp, cached)
    return cached


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 l["status"] == "ready":
            if l.get("model") not in MODELS:
                raise BuildError(f"languages.json: {code} is ready with an unknown model {l.get('model')!r}")
            ex = l.get("examples", {})
            if not ex.get("sentence") or any(len(s["out"]) != 3 for s in ex["sentence"]):
                raise BuildError(f"languages.json: {code} needs sentence examples of 3")
            if not ex.get("reply") or any(len(r["out"]) != 4 for r in ex["reply"]):
                raise BuildError(f"languages.json: {code} needs reply examples of 4")
        elif l["status"] != "not-yet":
            raise BuildError(f"languages.json: {code} status {l['status']!r}")
    return dict(sorted(langs.items()))


def split(model_id, source):
    """The model file in parts: (files to write, the parts as the catalogue lists them)."""
    files, parts = {}, []
    with open(source, "rb") as f:
        n = 0
        while True:
            data = f.read(PART_BYTES)
            if not data:
                break
            digest = hashlib.sha256(data).hexdigest()
            rel = f"{model_id}/part-{n:02d}.{digest[:12]}.bin"
            files[rel] = ("part", source, n * PART_BYTES, len(data), digest)
            parts.append({"url": rel, "bytes": len(data), "sha256": digest})
            n += 1
    if sum(p["bytes"] for p in parts) != MODELS[model_id]["bytes"]:
        raise BuildError(f"the parts of {model_id} do not add up to the model")
    return files, parts


def plan(sources, model_ids):
    """Every file to write, as rel -> (bytes or (src, offset, size)), and the catalogue."""
    if sha256_file(APACHE) != APACHE_SHA256:
        raise BuildError(f"{APACHE} is not the Apache License 2.0 text pinned here")
    apache = APACHE.read_bytes()
    apache_rel = f"licences/Apache-2.0.{hashlib.sha256(apache).hexdigest()[:12]}.txt"
    files = {apache_rel: ("data", apache)}
    models = {}
    for model_id in model_ids:
        m = MODELS[model_id]
        part_files, parts = split(model_id, sources[model_id])
        files.update(part_files)
        models[model_id] = {
            "name": m["name"],
            "runtime": RUNTIME,
            "source": f"https://huggingface.co/{m['repo']}/blob/{m['revision']}/{m['file']}",
            "bytes": m["bytes"],
            "sha256": m["sha256"],
            "parts": parts,
            "licence": LICENCE,
            "licence_url": LICENCE_URL,
            "attribution": ATTRIBUTION.format(name=m["name"]),
            "licence_file": apache_rel,
            "gpu_memory_bytes": m["gpu_memory_bytes"],
            "context_tokens": CONTEXT_TOKENS,
            "min_memory_gb": m["min_memory_gb"],
            "covers": [c for c in m["covers"] if c in model_ids],
        }
    catalogue = {"schema": SCHEMA, "generated": RETRIEVED, "models": models, "languages": languages()}
    cat_bytes = dumps(catalogue)
    cat_name = f"index.{hashlib.sha256(cat_bytes).hexdigest()[: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, apache_rel).encode("utf-8"))
    return files, cat_bytes


def licences_md(catalogue, apache_rel):
    out = [
        "# Sentence helper: licences\n\n"
        "The sentence helper in AuraSense AAC (aac-web) downloads one language model, on request, and runs it on "
        "the person's device. Nothing they type is sent anywhere.\n"
    ]
    for model_id, m in catalogue["models"].items():
        ready = [f"{l['name']} ({c})" for c, l in catalogue["languages"].items()
                 if l["status"] == "ready" and l.get("model") == model_id]
        out.append(
            f"\n## {m['name']}\n\n"
            f"- {m['attribution']}.\n"
            f"- Licence: {m['licence']} ({m['licence_url']}); the text is `{apache_rel}`.\n"
            f"- Source: {m['source']} ({m['bytes']} bytes, SHA-256 `{m['sha256']}`), unchanged, published here in "
            f"{len(m['parts'])} parts of at most 64 MiB, each named by its own SHA-256. Together they are the file.\n"
            f"- Size: {m['bytes'] / 1e9:.1f} GB. Ready in: {', '.join(ready) or 'no language yet'}.\n"
        )
    out.append(
        "\nThe runtime that runs it (LiteRT-LM, Apache-2.0) is vendored with its own licence and notice under "
        "`vendor/litert-lm-<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)

    # Parts 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)

    for rel in sorted(files, key=order):
        spec = files[rel]
        dest = out / rel
        dest.parent.mkdir(parents=True, exist_ok=True)
        if spec[0] == "part":
            _, src, offset, size, digest = spec
            fresh = dest.exists() and dest.stat().st_size == size and sha256_file(dest) == digest
            if not fresh:
                tmp = dest.with_name(dest.name + ".tmp")
                with open(src, "rb") as f, open(tmp, "wb") as o:
                    f.seek(offset)
                    o.write(f.read(size))
                os.replace(tmp, dest)
        else:
            data = spec[1]
            if not (dest.exists() and dest.read_bytes() == data):
                tmp = dest.with_name(dest.name + ".tmp")
                tmp.write_bytes(data)
                os.replace(tmp, dest)
        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", action="append", default=[], metavar="ID=FILE",
                    help="a local copy of a pinned model file")
    ap.add_argument("--offline", action="store_true", help="never download")
    ap.add_argument("--models", help="comma-separated model ids (default: those a ready language uses)")
    ap.add_argument("--out", type=Path, default=OUT_DEFAULT)
    ap.add_argument("--no-commit", action="store_true", help="leave data/ai alone (an evaluation build)")
    args = ap.parse_args()
    try:
        langs = languages()
        model_ids = (args.models.split(",") if args.models
                     else sorted({l["model"] for l in langs.values() if l["status"] == "ready"}))
        unknown = [m for m in model_ids if m not in MODELS]
        if unknown or not model_ids:
            raise BuildError(f"models {unknown or model_ids}: choose from {', '.join(MODELS)}")
        # A published build has every model a ready language uses; one with
        # fewer is an evaluation build, written elsewhere and not committed.
        used = sorted({l["model"] for l in langs.values() if l["status"] == "ready"})
        left_out = [m for m in used if m not in model_ids]
        if left_out and (not args.no_commit or args.out == OUT_DEFAULT):
            raise BuildError(f"ready languages use {', '.join(left_out)}: build every model they use, "
                             "or make an evaluation build with --out DIR --no-commit")
        given = {}
        for s in args.source:
            model_id, sep, file = s.partition("=")
            if not sep or model_id not in model_ids:
                raise BuildError(f"--source {s!r}: give ID=FILE, with ID one of {', '.join(model_ids)}")
            given[model_id] = file
        sources = {m: model_source(MODELS[m], given.get(m), args.offline) for m in model_ids}
        files, cat_bytes = plan(sources, model_ids)
        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
    parts = sum(1 for s in files.values() if s[0] == "part")
    log(f"✓ {args.out}: {', '.join(model_ids)}, {parts} parts, catalogue {json.loads(files['index.json'][1])['current']}")
    return 0


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