#!/usr/bin/env python3
"""The module corpus of the JPEG AI exact decoder module (waism §9.2), from committed files.

    python3 -I corpus.py           write corpus.json beside this file
    python3 -I corpus.py --check   require corpus.json to be what this writes, byte for byte

corpus.json holds the cases the module is held to: the capability, the payload, the set,
and what the module gives, the SHA-256 of its output or its refusal code. host/ reads it,
gives the module each case as waism §5 states, and holds every output to it. Nothing here is
new data: every payload, set name and digest is a committed file, or read from one.

  wai-rs/tests/vectors/jpegai/test256.bin   the reference codestream, the payload
  tools/jpegai-e2e-v1.sha256                the jpegai/v1 bundle: each file's name and
                                            SHA-256, fetched by
                                            tools/wai_jpegai_fetch_e2e_inputs.sh, which
                                            checks each against this manifest
  wai-rs/tests/jpegai_cross_platform.rs     PINNED's `rgb`: the SHA-256 of the RGB8 the
                                            native exact decoder gives for that bundle and
                                            codestream, the digest the vector expects
  wai-rs/tests/vectors/jpegai/configs/      the configuration vectors and their manifest
                                            (tools/wai_jpegai_configs.py): each one in
                                            scope with the RGB8 digest the native exact
                                            decoder gives, each other with its refusal

The bundle is not in this repository: the cases name its files under "external", and a
runner is given the directory the fetch script wrote them to. In the set, each file has the
name the set's listing gives it (SPEC §3.1): `norm.json`, and each weight file under
`weights/`; the set is their file sequence, in ascending byte order of the names (waism §5).

The refusal fixtures are the bundle with one member changed, each refused with the code the
capability's text gives (SPEC §5, the `wai.image.jpegai` row): a member missing, and a
member that is not JSON, the first read (`norm.json`) and the last (`weights/icci.json`).

The configuration cases (`jpegai/configs/<name>`) are the configuration vectors over the same
bundle: a vector in scope expects the native exact decoder's digest from the manifest, a
vector of a configuration the reference software leaves to the machine expects the refusal
`configuration_undetermined`, and a vector the decoder does not implement expects the sink's
refusal `module_trap` (the module traps; waism §6 check 12).
"""

import hashlib
import json
import os
import re
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
WAI = os.path.dirname(HERE)
OUT = os.path.join(HERE, "corpus.json")

CAPABILITY = "wai.image.jpegai"
PAYLOAD = "wai-rs/tests/vectors/jpegai/test256.bin"
MANIFEST = "tools/jpegai-e2e-v1.sha256"
PINNED = "wai-rs/tests/jpegai_cross_platform.rs"
CONFIGS = "wai-rs/tests/vectors/jpegai/configs"
FETCH = "tools/wai_jpegai_fetch_e2e_inputs.sh"


def read(rel):
    with open(os.path.join(WAI, rel), "rb") as f:
        return f.read()


def manifest():
    """The bundle's files, name → SHA-256, from the manifest the fetch script checks."""
    out = {}
    for line in read(MANIFEST).decode("utf-8").splitlines():
        if line.strip():
            sha, name = line.split()
            out[name] = sha
    return out


def pinned_rgb():
    """PINNED's `rgb`: the native exact decoder's RGB8 digest for the jpegai/v1 inputs."""
    text = read(PINNED).decode("utf-8")
    block = re.search(r"const PINNED: \[\(&str, &str\); \d+\] = \[(.*?)\];", text, re.S)
    if not block:
        sys.exit(f"corpus.py: {PINNED} has no PINNED table")
    rgb = re.findall(r'\("rgb", "([0-9a-f]{64})"\)', block.group(1))
    if len(rgb) != 1:
        sys.exit(f"corpus.py: {PINNED}'s PINNED table has no one rgb digest")
    return rgb[0]


def corpus():
    files = manifest()
    payload_sha = files.pop("test256.bin", None)
    if payload_sha != hashlib.sha256(read(PAYLOAD)).hexdigest():
        sys.exit(f"corpus.py: {PAYLOAD} is not the test256.bin of {MANIFEST}")
    # Each member by its name in the set's listing: norm.json, and the weights under weights/.
    members = sorted(
        ("norm.json" if n == "norm.json" else f"weights/{n}", {"external": n}) for n in files
    )
    names = [m[0] for m in members]
    if len(names) != 12 or names != sorted(names, key=lambda s: s.encode("utf-8")):
        sys.exit("corpus.py: the bundle is not its twelve members in byte order")

    def without(name):
        return [list(m) for m in members if m[0] != name]

    def changed(name, hexbytes):
        return [[m[0], {"hex": hexbytes} if m[0] == name else m[1]] for m in members]

    cases = [
        {"name": "jpegai/v1/test256", "capability": CAPABILITY, "payload": {"file": PAYLOAD},
         "set": [list(m) for m in members], "expect": {"sha256": pinned_rgb()}},
        {"name": "jpegai/v1/refuse/icci-missing", "capability": CAPABILITY,
         "payload": {"file": PAYLOAD}, "set": without("weights/icci.json"),
         "expect": {"refuse": "set_incomplete"}},
        {"name": "jpegai/v1/refuse/norm-not-json", "capability": CAPABILITY,
         "payload": {"file": PAYLOAD}, "set": changed("norm.json", "7b"),
         "expect": {"refuse": "model_mismatch"}},
        {"name": "jpegai/v1/refuse/icci-not-json", "capability": CAPABILITY,
         "payload": {"file": PAYLOAD}, "set": changed("weights/icci.json", "7b"),
         "expect": {"refuse": "model_mismatch"}},
    ]
    m = json.loads(read(CONFIGS + "/manifest.json"))
    for v in m["vectors"]:
        if v["name"] == "test256":
            continue
        path = os.path.normpath(os.path.join(CONFIGS, v["file"]))
        if hashlib.sha256(read(path)).hexdigest() != v["sha256"]:
            sys.exit(f"corpus.py: {path} is not the manifest's")
        e = v["expect"]
        if "rgb_sha256" in e:
            expect = {"sha256": e["rgb_sha256"]}
        elif e["refuse"] == "configuration_undetermined":
            expect = {"refuse": "configuration_undetermined"}
        else:
            expect = {"sink": "module_trap"}
        cases.append({"name": "jpegai/configs/" + v["name"], "capability": CAPABILITY,
                      "payload": {"file": path}, "set": [list(m_) for m_ in members], "expect": expect})
    doc = {
        "format": "waism-module-corpus/1",
        "about": (
            "The module corpus (extensions/waism.md §9.2) of wai-det-jpegai, the JPEG AI exact "
            "decoder at WAISM's float level: for each case, the capability, the payload and the "
            "set the module is given (waism §5), and the SHA-256 of its output or its refusal "
            "code, or the sink's refusal of the run. The vectors' digests are the native exact "
            f"decoder's pinned RGB8 ({PINNED}, PINNED's rgb; {CONFIGS}/manifest.json). Written by "
            "corpus.py; paths are below wai/."
        ),
        "external": {n: {"sha256": s, "fetch": FETCH} for n, s in sorted(files.items())},
        "cases": cases,
    }
    return (json.dumps(doc, indent=1, ensure_ascii=False) + "\n").encode("utf-8")


def main(argv):
    want = corpus()
    if argv[1:] == ["--check"]:
        have = open(OUT, "rb").read() if os.path.isfile(OUT) else b""
        if have != want:
            print("corpus.py --check: corpus.json is not what corpus.py writes", file=sys.stderr)
            return 1
        print("corpus.py --check: corpus.json is what corpus.py writes")
        return 0
    if argv[1:]:
        print("usage: corpus.py [--check]", file=sys.stderr)
        return 2
    with open(OUT, "wb") as f:
        f.write(want)
    print(f"corpus.py: wrote {OUT}")
    return 0


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