#!/usr/bin/env python3
"""Independent check of the pin-conformance corpus (SPEC.md §3.1).

Shares no code with the Rust reference implementation: it re-implements the
listing rules and the ordered pin checks from the SPEC text, with the standard
library only, and checks that

  - every `listings/<name>.sha256` is byte for byte the listing of `sets/<name>/`
    (regular files, names that begin with `.` skipped), and every file under
    `listings/bad/` breaks a listing rule;
  - every case in `cases.json` gets its declared reason.

  python3 pin-conformance/pin_verify.py [corpus-dir]

The registry facts it needs are restated here from SPEC.md §5 and §3.1: the
registered prior forms, the forms the reference sink reads, and the capabilities
registered with a one-file parameter set.
"""
import hashlib
import json
import os
import re
import sys

PRIOR_FORMS = {'wai.prior.nnc', 'wai.prior.set'}
READABLE = ['wai.prior.set']
ONE_FILE = {
    'wai.neural.int_synth', 'wai.neural.int_mlicpp', 'wai.neural.encodec32', 'wai.neural.dac',
    'wai.neural.mimi', 'wai.neural.wavtokenizer', 'wai.neural.snac', 'wai.neural.video_bmshj2018',
}
HEX64 = re.compile(r'^[0-9a-f]{64}$')
SEGMENT = re.compile(r'^[A-Za-z0-9_-]([A-Za-z0-9_.-]*[A-Za-z0-9_-])?$')


class Malformed(Exception):
    pass


def sha(b):
    return hashlib.sha256(b).hexdigest()


def parse_listing(b):
    if not b or not b.endswith(b'\n'):
        raise Malformed('empty, or no final line feed')
    lines = []
    for raw in b[:-1].split(b'\n'):
        if len(raw) < 67 or raw[64:66] != b'  ':
            raise Malformed('line shape')
        d = raw[:64].decode('ascii', 'replace')
        if not HEX64.match(d):
            raise Malformed('digest')
        try:
            p = raw[66:].decode('ascii')
        except UnicodeDecodeError:
            raise Malformed('path not ASCII')
        if not all(SEGMENT.match(s) for s in p.split('/')):
            raise Malformed(f'path {p!r}')
        lines.append((p, d))
    paths = [p for p, _ in lines]
    if any(a.encode() >= b.encode() for a, b in zip(paths, paths[1:])):
        raise Malformed('order')
    folded = sorted(p.lower() for p in paths)
    if any(a == b for a, b in zip(folded, folded[1:])):
        raise Malformed('case fold')
    for a in paths:
        if any(b.startswith(a + '/') for b in paths):
            raise Malformed('file and directory')
    return lines


def listing_of_dir(root):
    files = []
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = [d for d in dirnames if not d.startswith('.')]
        rel = os.path.relpath(dirpath, root)
        for f in filenames:
            if f.startswith('.'):
                continue
            if rel == '.' and f.endswith('.sha256'):
                continue
            p = f if rel == '.' else rel.replace(os.sep, '/') + '/' + f
            files.append(p)
    out = b''
    for p in sorted(files, key=lambda x: x.encode()):
        with open(os.path.join(root, p), 'rb') as fh:
            out += sha(fh.read()).encode() + b'  ' + p.encode() + b'\n'
    return out


def verify(capability, pin, held, readable):
    """SPEC §3.1's checks, in order; returns 'ok' or the reason."""
    if not HEX64.match(pin['sha256']):
        return 'pin_malformed'
    form = pin.get('form')
    if form is not None:
        if form not in PRIOR_FORMS:
            return 'form_not_registered'
        if form not in readable:
            return 'form_unreadable'
    elif capability not in ONE_FILE:
        return 'form_required'
    set_form = form == 'wai.prior.set'
    digest = lambda h: sha(h['bytes']) if (('files' in h) == set_form) else None
    matching = [h for h in held if digest(h) == pin['sha256']]
    if not matching:
        sup = pin.get('supersedes')
        if (sup is not None and HEX64.match(sup['sha256']) and sup['sha256'] != pin['sha256']
                and (sup.get('form') is None or sup.get('form') in PRIOR_FORMS)):
            pre_set = sup.get('form') == 'wai.prior.set'
            if any(sha(h['bytes']) == sup['sha256'] and (('files' in h) == pre_set) for h in held):
                return 'digest_superseded'
        return 'prior_not_held'
    if not set_form:
        return 'ok'
    try:
        lines = parse_listing(matching[0]['bytes'])
    except Malformed:
        return 'set_malformed'
    for h in matching:
        if all(p in h['files'] and sha(h['files'][p]) == d for p, d in lines):
            return 'ok'
    return 'part_mismatch'


def load(corpus, held):
    out = []
    for h in held:
        if 'file' in h:
            with open(os.path.join(corpus, h['file']), 'rb') as f:
                out.append({'bytes': f.read()})
        else:
            with open(os.path.join(corpus, h['listing']), 'rb') as f:
                b = f.read()
            files = {}
            try:
                for p, _ in parse_listing(b):
                    fp = os.path.join(corpus, h['root'], p)
                    if os.path.isfile(fp):
                        with open(fp, 'rb') as f:
                            files[p] = f.read()
            except Malformed:
                pass
            out.append({'bytes': b, 'files': files})
    return out


def main():
    corpus = sys.argv[1] if len(sys.argv) > 1 else os.path.dirname(os.path.abspath(__file__))
    fails, n = [], 0
    ld = os.path.join(corpus, 'listings')
    for name in sorted(os.listdir(ld)):
        path = os.path.join(ld, name)
        if name.endswith('.sha256') and os.path.isfile(path):
            n += 1
            with open(path, 'rb') as f:
                committed = f.read()
            if listing_of_dir(os.path.join(corpus, 'sets', name[:-len('.sha256')])) != committed:
                fails.append(f'listings/{name} is not the listing of its set')
    for name in sorted(os.listdir(os.path.join(ld, 'bad'))):
        n += 1
        with open(os.path.join(ld, 'bad', name), 'rb') as f:
            b = f.read()
        try:
            parse_listing(b)
            fails.append(f'listings/bad/{name} parsed')
        except Malformed:
            pass
    with open(os.path.join(corpus, 'cases.json')) as f:
        cases = json.load(f)['cases']
    for c in cases:
        n += 1
        pin = {k: v for k, v in c['pin'].items() if v is not None}
        got = verify(c['capability'], pin, load(corpus, c['held']), c.get('readable', READABLE))
        if got != c['expect']:
            fails.append(f"cases.json {c['name']}: expected {c['expect']}, got {got}")
    for f in fails:
        print('FAIL', f)
    print(f'pin_verify.py: {n - len(fails)} of {n} checks as declared')
    sys.exit(1 if fails else 0)


if __name__ == '__main__':
    main()