#!/usr/bin/env python3
# Independent Falcon HashToPoint oracle using Python's standard-library SHAKE256
# (hashlib.shake_256 = FIPS 202). Cross-checks the AERE HashToPoint precompile
# (0x0AE7) output emitted by HashToPointKat.java.
#
# vectors file line: tcId|logn|nonce_hex|message_hex|precompile_coeffs_hex
# We recompute the challenge polynomial coefficients from (nonce||message) and
# assert they equal the precompile's coeffs column, byte-for-byte.
#
# Reference algorithm (Falcon hash_to_point_vartime): absorb nonce||message into
# SHAKE256, squeeze 2 bytes at a time as a big-endian uint16 w; keep w % q whenever
# w < 5*q = 61445, until n = 2^logn coefficients are collected.

import hashlib
import json
import sys

Q = 12289
REJECT = 5 * Q  # 61445


def hash_to_point(nonce: bytes, message: bytes, n: int):
    xof = hashlib.shake_256()
    xof.update(nonce + message)
    # squeeze a generous buffer; refill if a pathological vector needs more
    need = 2 * n
    buf = xof.digest(max(4096, need * 4))
    coeffs = []
    i = 0
    while len(coeffs) < n:
        if i + 1 >= len(buf):
            # extremely unlikely; extend the squeeze deterministically
            buf = xof.digest(len(buf) * 2)
        w = (buf[i] << 8) | buf[i + 1]
        i += 2
        if w < REJECT:
            coeffs.append(w % Q)
    return coeffs


def coeffs_to_hex(coeffs):
    out = bytearray()
    for c in coeffs:
        out.append((c >> 8) & 0xFF)
        out.append(c & 0xFF)
    return out.hex()


def main():
    vectors = sys.argv[1] if len(sys.argv) > 1 else "hashtopoint_vectors.txt"
    outpath = sys.argv[2] if len(sys.argv) > 2 else "kat-results-hashtopoint.json"

    passed = failed = total = 0
    results = []
    with open(vectors) as f:
        for raw in f:
            line = raw.strip()
            if not line or line.startswith("#"):
                continue
            p = line.split("|")
            if len(p) != 5:
                continue
            total += 1
            tc, logn, nonce_hex, msg_hex, coeffs_hex = p
            n = 1 << int(logn)
            nonce = bytes.fromhex(nonce_hex)
            message = bytes.fromhex(msg_hex)
            ref = coeffs_to_hex(hash_to_point(nonce, message, n))
            ok = (ref == coeffs_hex.lower())
            if ok:
                passed += 1
            else:
                failed += 1
            results.append({
                "tcId": tc, "logn": int(logn), "n": n,
                "nonceLen": len(nonce), "msgLen": len(message),
                "status": "PASS" if ok else "FAIL",
            })
            print(f"HashToPoint tc{tc} logn={logn} n={n} -> {'PASS' if ok else 'FAIL'}")

    report = {
        "precompile": "Falcon HashToPoint (SHAKE256 rejection sampler)",
        "address": "0x0000000000000000000000000000000000000ae7",
        "oracle": "python hashlib.shake_256 (FIPS 202), independent of Bouncy Castle",
        "passed": passed, "failed": failed, "total": total, "results": results,
    }
    with open(outpath, "w") as f:
        json.dump(report, f, indent=2)

    print(f"HashToPoint KAT: passed={passed} failed={failed} total={total}")
    if failed != 0 or total == 0:
        print("HASHTOPOINT_KAT_FAILED")
        sys.exit(3)
    print("HASHTOPOINT_KAT_OK")


if __name__ == "__main__":
    main()