#!/usr/bin/env python3
"""Write the set-form twin of the `wai.neural.int_mlicpp` conformance vector.

The committed `stream.wai` pins `model.wmb` form-less: its pin's `sha256` is the
SHA-256 of the one file. SPEC.md §3.1 also lets a pin cover a parameter set through
its listing (the prior form `wai.prior.set`). This tool writes, next to the vector:

  model.sha256    the listing of the one-file set: the line
                  `<sha256 of model.wmb>  model.wmb` and a line feed -- the text
                  `sha256sum model.wmb` prints, which `sha256sum -c` checks.
  stream_set.wai  `stream.wai` with the same payload, its pin replaced by one over
                  the listing: `sha256` = SHA-256(model.sha256), `form`
                  `wai.prior.set`, `tier` `decode-equivalence`.

Standard library only, and independent of the Rust reference implementation, which
checks these files in tests/int_mlicpp.rs and src/pin.rs.

  python tools/wai_int_mlicpp_set_form.py [vector-dir]

`vector-dir` defaults to wai-rs/tests/vectors/int_mlicpp. `wai_int_mlicpp_export.py`
calls `write_set_form` after it writes the vector.
"""
import hashlib
import json
import os
import struct
import sys


def listing(files):
    """The listing of `files` (path -> bytes): sorted by path as bytes, one
    `<64 hex>  <path>\\n` line per file."""
    out = b''
    for path in sorted(files, key=lambda p: p.encode()):
        out += hashlib.sha256(files[path]).hexdigest().encode() + b'  ' + path.encode() + b'\n'
    return out


def write_set_form(vector_dir):
    with open(os.path.join(vector_dir, 'model.wmb'), 'rb') as f:
        wmb = f.read()
    with open(os.path.join(vector_dir, 'stream.wai'), 'rb') as f:
        env = f.read()
    assert env[:4] == b'WAI1'
    ml = struct.unpack('<I', env[4:8])[0]
    manifest = json.loads(env[8:8 + ml])
    rest = env[8 + ml:]
    lst = listing({'model.wmb': wmb})
    prior = manifest['conditioning']['prior']
    assert prior['sha256'] == hashlib.sha256(wmb).hexdigest(), 'stream.wai pins model.wmb'
    prior['sha256'] = hashlib.sha256(lst).hexdigest()
    prior['form'] = 'wai.prior.set'
    prior['tier'] = 'decode-equivalence'
    mb = json.dumps(manifest, separators=(',', ':')).encode()
    set_env = b'WAI1' + struct.pack('<I', len(mb)) + mb + rest
    with open(os.path.join(vector_dir, 'model.sha256'), 'wb') as f:
        f.write(lst)
    with open(os.path.join(vector_dir, 'stream_set.wai'), 'wb') as f:
        f.write(set_env)
    return lst, set_env


def main():
    here = os.path.dirname(os.path.abspath(__file__))
    d = sys.argv[1] if len(sys.argv) > 1 else os.path.join(here, '..', 'wai-rs', 'tests', 'vectors', 'int_mlicpp')
    lst, set_env = write_set_form(d)
    print(f'model.sha256 {len(lst)} B sha256 {hashlib.sha256(lst).hexdigest()}; '
          f'stream_set.wai {len(set_env)} B sha256 {hashlib.sha256(set_env).hexdigest()}')


if __name__ == '__main__':
    main()