#!/usr/bin/env python3
"""weft_import.py: copies the frames of the WeftCodec format's frozen conformance vectors
(conformance/v1/codec/ of the format's reference implementation, Apache-2.0) into
tensor-conformance/vectors/, one directory a vector holding frame.wft, and prints the rows
of the third-party table of wai-rs/src/codecs/weft/corpus.rs: each vector's name, the
SHA-256 of its frame, what the format's own vectors expect of it, and what a sink of
wai.tensor.weft gives it.

    weft_import.py <codec-vector-dir> <tensor-conformance-dir>

The vectors a later vector supersedes (their mode value has since been allocated) and the
one-mebibyte RAW vector are not copied. A frame of a mode wai.tensor.weft admits (0, 6, 9)
keeps the format's expectation; any other frame whose header is sound is UnsupportedMode,
the class the format gives a mode a decoder does not implement. Standard library only.
"""
import hashlib
import json
import os
import struct
import sys

ADMITTED = {0, 6, 9}
SKIP = {"004_raw_1m"}


def crc32c(data):
    c = 0xFFFFFFFF
    for b in data:
        c ^= b
        for _ in range(8):
            c = (c >> 1) ^ 0x82F63B78 if c & 1 else c >> 1
    return c ^ 0xFFFFFFFF


def header_sound(frame):
    """Magic, version and header CRC as a decoder checks them before reading the mode."""
    if len(frame) < 100 or frame[:4] != b"WFT1" or frame[4] != 1 or frame[5] != 0:
        return False
    return crc32c(frame[:72] + frame[76:96]) == struct.unpack_from("<I", frame, 72)[0]


def main():
    src, dst = sys.argv[1], sys.argv[2]
    rows = []
    for fn in sorted(os.listdir(src)):
        if not fn.endswith(".json"):
            continue
        with open(os.path.join(src, fn), encoding="utf-8") as fh:
            v = json.load(fh)
        name = v["name"]
        if name in SKIP or "superseded_by" in v:
            continue
        frame = bytes.fromhex(v["frame_hex"])
        d = os.path.join(dst, "vectors", name)
        os.makedirs(d, exist_ok=True)
        with open(os.path.join(d, "frame.wft"), "wb") as fh:
            fh.write(frame)
        exp = v["expect"]
        upstream = "ok" if exp.get("ok") else exp["error"]
        mode = frame[7] if len(frame) > 7 else None
        if header_sound(frame) and mode not in ADMITTED and mode is not None and mode <= 18:
            ours = "UnsupportedMode"
        elif header_sound(frame) and mode is not None and mode > 18:
            ours = "UnsupportedMode"
        else:
            ours = upstream
        # A limit belongs to the mode the vector exercises; a frame refused for its mode keeps none.
        limits = v.get("limits", {}) if ours != "UnsupportedMode" else {}
        plain = ""
        if ours == "ok":
            p = bytes.fromhex(v["plaintext_hex"])
            plain = f"Some(({len(p)}, \"{hashlib.sha256(p).hexdigest()}\"))"
        else:
            plain = "None"
        lim = ", ".join(f"(\"{k}\", {n})" for k, n in sorted(limits.items()))
        rows.append(f"    Upstream {{ name: \"{name}\", frame_sha256: \"{hashlib.sha256(frame).hexdigest()}\", "
                    f"upstream: \"{upstream}\", expect: \"{ours if ours != 'ok' else 'decode'}\", plaintext: {plain}, limits: &[{lim}] }},")
    print("\n".join(rows))


if __name__ == "__main__":
    main()