"""test_roundtrip.py — verify shake256-fanout matches hashlib.shake_256.

Reads the same input portal the cuda binary consumed, hashes each
input with Python's reference SHAKE256, compares against the binary's
output portal. If every hash matches, the cuda port is byte-identical
to upstream SHAKE256."""

import hashlib
import re
import sys


def extract_strings(text, section):
    """Pull "..."-quoted strings inside (section ...)."""
    m = re.search(r"\(" + section + r"\b([^)]*(?:\)[^)]+)*)", text, re.DOTALL)
    if not m:
        return []
    return re.findall(r'"([^"]*)"', m.group(1))


def main():
    in_path, out_path = sys.argv[1], sys.argv[2]
    with open(in_path) as fh:
        in_text = fh.read()
    with open(out_path) as fh:
        out_text = fh.read()

    m = re.search(r"\(output-bytes\s+(\d+)\)", in_text)
    out_bytes = int(m.group(1)) if m else 32

    inputs = extract_strings(in_text, "inputs")
    hashes = extract_strings(out_text, "hashes")

    if len(inputs) != len(hashes):
        print(f"COUNT MISMATCH: {len(inputs)} inputs vs {len(hashes)} hashes")
        sys.exit(1)

    n_ok = 0
    for i, (hex_in, got_hex) in enumerate(zip(inputs, hashes)):
        expected = hashlib.shake_256(bytes.fromhex(hex_in)).hexdigest(out_bytes)
        if expected != got_hex:
            print(f"FAIL #{i}: input={hex_in[:20]}...")
            print(f"  expected {expected}")
            print(f"  got      {got_hex}")
            sys.exit(1)
        n_ok += 1
    print(f"PASS — {n_ok} / {n_ok} hashes byte-identical to hashlib.shake_256")


if __name__ == "__main__":
    main()
