#!/usr/bin/env python3
"""Check a first-letters directory before a site build ships it.

  python3 tools/ae/check_packs.py DIR

DIR must hold what tools/ae/build_ae.py built for the committed catalogue
(data/ae/index.json): its index.json points at index.<sha12>.json, whose
bytes are the committed catalogue's, and every part of the graph and its
weights, the tokenizer, the configuration, the letter tables and every
licence that catalogue names is there with its exact size and SHA-256.
Otherwise this exits non-zero and says what is wrong, so a build never ships
a model that disagrees with the catalogue, or that is incomplete.

Standard library only; Python 3.9 or later (macOS's own python3).
"""

import hashlib
import json
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
COMMITTED = HERE.parents[1] / "data" / "ae" / "index.json"
SHOW = 10


def sha256(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 model_files(model):
    """Every file a model names: its graph, its weights' parts, its tokenizer,
    configuration and letter tables."""
    out = [model["graph"]] if "graph" in model else []
    for d in model.get("data", []):
        out += d.get("parts", [])
    for key in ("tokenizer", "tokenizer_config", "config", "vocab", "vocab_table"):
        if key in model:
            out.append(model[key])
    return [(f["url"], f["bytes"], f["sha256"]) for f in out]


def main(argv):
    if len(argv) != 2:
        print(__doc__, file=sys.stderr)
        return 2
    root = Path(argv[1])
    problems = []
    try:
        pointer = json.loads((root / "index.json").read_text(encoding="utf-8"))
        name = pointer["current"]
        cat_bytes = (root / name).read_bytes()
    except (OSError, ValueError, KeyError) as e:
        print(f"✗ {root}: no readable index.json and catalogue ({e})", file=sys.stderr)
        return 1
    if cat_bytes != COMMITTED.read_bytes():
        problems.append(f"{name} is not the committed catalogue ({COMMITTED})")
    if name != f"index.{hashlib.sha256(cat_bytes).hexdigest()[:12]}.json":
        problems.append(f"{name} is not named by its contents")
    cat = json.loads(cat_bytes)
    refs = []
    for model in cat.get("models", {}).values():
        refs += model_files(model)
        for credit in model.get("credits", []):
            lf = credit.get("licence_file")
            if lf:
                data = (root / lf).read_bytes() if (root / lf).is_file() else b""
                refs.append((lf, len(data) if data else -1, hashlib.sha256(data).hexdigest() if data else ""))
    if not (root / "LICENSES.md").is_file():
        problems.append("LICENSES.md is missing")
    for rel, size, digest in refs:
        f = root / rel
        if not f.is_file():
            problems.append(f"missing {rel}")
        elif f.stat().st_size != size:
            problems.append(f"{rel} is {f.stat().st_size} bytes, the catalogue says {size}")
        elif sha256(f) != digest:
            problems.append(f"{rel} does not match its SHA-256")
    if problems:
        for p in problems[:SHOW]:
            print(f"✗ {p}", file=sys.stderr)
        if len(problems) > SHOW:
            print(f"  … and {len(problems) - SHOW} more", file=sys.stderr)
        return 1
    print(f"✓ {root}: {len(refs)} files match {name}")
    return 0


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