// SPDX-License-Identifier: MIT
pragma solidity 0.8.23;

/**
 * @title AereFalcon512Verifier, spec-complete NIST Falcon-512 signature verifier, on-chain
 * @notice A REAL, standalone Falcon-512 (FIPS 206 / NIST PQC round-3) signature
 *         verifier that runs entirely on chain 2800. Given a public key, a
 *         message, and a signature, it performs the full Falcon verification:
 *
 *           1. HashToPoint: SHAKE256(nonce || message) is streamed and
 *              rejection-sampled into a challenge polynomial c in Z_q[x]/(x^512+1)
 *              (q = 12289), exactly per the reference hash_to_point_vartime.
 *           2. Public-key decode: the 897-byte NIST public key (0x09 header + 512
 *              coefficients packed at 14 bits each) is decoded to h.
 *           3. Signature decode: the compressed Gaussian encoding of s2 is decoded
 *              (Falcon comp_decode: sign bit, 7 low bits, unary high bits).
 *           4. Recompute -s1 = s2*h - c in the ring (negacyclic convolution mod q),
 *              then center coefficients into (-q/2, q/2].
 *           5. Norm check: accept iff ||(s1, s2)||^2 <= floor(beta^2) = 34034726,
 *              the Falcon-512 acceptance bound.
 *
 *         This is the actual lattice/hash cryptography, executed on-chain, not a
 *         proof-of-a-proof. It is validated bit-for-bit against the OFFICIAL NIST
 *         Falcon-512 Known-Answer-Test vectors (falcon512-KAT.rsp, SHA-256
 *         dd75c946...b5cd, reproduced byte-for-byte from the round-3 reference):
 *         real (pk, message, signature) triples verify true; tampered signatures
 *         verify false.
 *
 *         The chain's block gas limit is effectively unlimited, so the full
 *         verification runs as an ordinary transaction; there is no trusted
 *         prover and no off-chain step. Anyone can call verify() and get the same
 *         accept/reject the reference C implementation produces.
 */
contract AereFalcon512Verifier {
    // ----- Falcon-512 parameters -----
    uint256 public constant Q = 12289; // modulus
    uint256 public constant N = 512; // ring degree
    uint256 public constant L2BOUND = 34034726; // floor(beta^2), Falcon-512 acceptance bound
    uint256 internal constant HQ = 6144; // floor(q/2), centering threshold

    // Negacyclic NTT constants for Z_q[x]/(x^512+1), q = 12289:
    //   PSI = primitive 1024th (2n-th) root of unity; W = PSI^2 (512th root)
    uint256 internal constant PSI = 49;
    uint256 internal constant PSI_INV = 1254;
    uint256 internal constant W = 2401; // PSI^2 mod q
    uint256 internal constant W_INV = 11813;
    uint256 internal constant N_INV = 12265; // 512^-1 mod q
    uint8 public constant PK_HEADER = 0x09; // 0x00 + logn, logn = 9
    uint8 public constant SIG_HEADER = 0x29; // 0x20 + logn, logn = 9 (NIST esig header)
    uint256 public constant NONCE_LEN = 40;

    // ----- recorded on-chain verification results (for state-changing calls) -----
    uint256 public verifyCount;
    bool public lastResult;
    event Verified(address indexed caller, bool result, uint256 index);

    // ==========================================================================
    //                       Keccak-f[1600] / SHAKE256
    // ==========================================================================

    /// @dev In-place Keccak-f[1600] permutation on a 25-lane state (one lane per
    ///      32-byte memory word, low 64 bits used). Written in assembly for gas:
    ///      the naive array version cost ~1.26M gas per permutation. Validated
    ///      against the FIPS 202 SHAKE256 KAT and the Falcon HashToPoint oracle.
    function _keccakf(uint64[25] memory st) internal pure {
        assembly {
            function rol(x, n) -> r {
                r := and(0xffffffffffffffff, or(shl(n, x), shr(sub(64, n), x)))
            }
            let s := st
            let b := mload(0x40) // scratch (25 lanes), left above the free pointer
            let rc := add(b, 0x320) // 24 round constants after the scratch lanes
            // round constants
            mstore(add(rc, 0x000), 0x0000000000000001) mstore(add(rc, 0x020), 0x0000000000008082)
            mstore(add(rc, 0x040), 0x800000000000808a) mstore(add(rc, 0x060), 0x8000000080008000)
            mstore(add(rc, 0x080), 0x000000000000808b) mstore(add(rc, 0x0a0), 0x0000000080000001)
            mstore(add(rc, 0x0c0), 0x8000000080008081) mstore(add(rc, 0x0e0), 0x8000000000008009)
            mstore(add(rc, 0x100), 0x000000000000008a) mstore(add(rc, 0x120), 0x0000000000000088)
            mstore(add(rc, 0x140), 0x0000000080008009) mstore(add(rc, 0x160), 0x000000008000000a)
            mstore(add(rc, 0x180), 0x000000008000808b) mstore(add(rc, 0x1a0), 0x800000000000008b)
            mstore(add(rc, 0x1c0), 0x8000000000008089) mstore(add(rc, 0x1e0), 0x8000000000008003)
            mstore(add(rc, 0x200), 0x8000000000008002) mstore(add(rc, 0x220), 0x8000000000000080)
            mstore(add(rc, 0x240), 0x000000000000800a) mstore(add(rc, 0x260), 0x800000008000000a)
            mstore(add(rc, 0x280), 0x8000000080008081) mstore(add(rc, 0x2a0), 0x8000000000008080)
            mstore(add(rc, 0x2c0), 0x0000000080000001) mstore(add(rc, 0x2e0), 0x8000000080008008)

            for { let rnd := 0 } lt(rnd, 24) { rnd := add(rnd, 1) } {
                // ---- theta: column parities C[x] -> b[0..4] ----
                for { let x := 0 } lt(x, 5) { x := add(x, 1) } {
                    let o := mul(x, 0x20)
                    mstore(add(b, o),
                        xor(xor(xor(xor(
                            mload(add(s, o)),
                            mload(add(s, add(o, 0xa0)))),
                            mload(add(s, add(o, 0x140)))),
                            mload(add(s, add(o, 0x1e0)))),
                            mload(add(s, add(o, 0x280)))))
                }
                // D[x] = C[(x+4)%5] ^ rol(C[(x+1)%5],1) -> b[5..9]
                for { let x := 0 } lt(x, 5) { x := add(x, 1) } {
                    let cm := mload(add(b, mul(mod(add(x, 4), 5), 0x20)))
                    let cp := mload(add(b, mul(mod(add(x, 1), 5), 0x20)))
                    mstore(add(b, add(0xa0, mul(x, 0x20))), xor(cm, rol(cp, 1)))
                }
                // apply D to state: lane (x + 5y) ^= D[x]
                for { let x := 0 } lt(x, 5) { x := add(x, 1) } {
                    let d := mload(add(b, add(0xa0, mul(x, 0x20))))
                    let o := mul(x, 0x20)
                    mstore(add(s, o), xor(mload(add(s, o)), d))
                    mstore(add(s, add(o, 0xa0)), xor(mload(add(s, add(o, 0xa0))), d))
                    mstore(add(s, add(o, 0x140)), xor(mload(add(s, add(o, 0x140))), d))
                    mstore(add(s, add(o, 0x1e0)), xor(mload(add(s, add(o, 0x1e0))), d))
                    mstore(add(s, add(o, 0x280)), xor(mload(add(s, add(o, 0x280))), d))
                }
                // ---- rho + pi: b[piDest[i]] = rol(st[i], rho[i]) ----
                mstore(add(b, 0x000), rol(mload(add(s, 0x000)), 0))
                mstore(add(b, 0x140), rol(mload(add(s, 0x020)), 1))
                mstore(add(b, 0x280), rol(mload(add(s, 0x040)), 62))
                mstore(add(b, 0x0a0), rol(mload(add(s, 0x060)), 28))
                mstore(add(b, 0x1e0), rol(mload(add(s, 0x080)), 27))
                mstore(add(b, 0x200), rol(mload(add(s, 0x0a0)), 36))
                mstore(add(b, 0x020), rol(mload(add(s, 0x0c0)), 44))
                mstore(add(b, 0x160), rol(mload(add(s, 0x0e0)), 6))
                mstore(add(b, 0x2a0), rol(mload(add(s, 0x100)), 55))
                mstore(add(b, 0x0c0), rol(mload(add(s, 0x120)), 20))
                mstore(add(b, 0x0e0), rol(mload(add(s, 0x140)), 3))
                mstore(add(b, 0x220), rol(mload(add(s, 0x160)), 10))
                mstore(add(b, 0x040), rol(mload(add(s, 0x180)), 43))
                mstore(add(b, 0x180), rol(mload(add(s, 0x1a0)), 25))
                mstore(add(b, 0x2c0), rol(mload(add(s, 0x1c0)), 39))
                mstore(add(b, 0x2e0), rol(mload(add(s, 0x1e0)), 41))
                mstore(add(b, 0x100), rol(mload(add(s, 0x200)), 45))
                mstore(add(b, 0x240), rol(mload(add(s, 0x220)), 15))
                mstore(add(b, 0x060), rol(mload(add(s, 0x240)), 21))
                mstore(add(b, 0x1a0), rol(mload(add(s, 0x260)), 8))
                mstore(add(b, 0x1c0), rol(mload(add(s, 0x280)), 18))
                mstore(add(b, 0x300), rol(mload(add(s, 0x2a0)), 2))
                mstore(add(b, 0x120), rol(mload(add(s, 0x2c0)), 61))
                mstore(add(b, 0x260), rol(mload(add(s, 0x2e0)), 56))
                mstore(add(b, 0x080), rol(mload(add(s, 0x300)), 14))
                // ---- chi: st[y+x] = b[y+x] ^ ((~b[y+(x+1)%5]) & b[y+(x+2)%5]) ----
                for { let yo := 0 } lt(yo, 0x320) { yo := add(yo, 0xa0) } {
                    let b0 := mload(add(b, yo))
                    let b1 := mload(add(b, add(yo, 0x20)))
                    let b2 := mload(add(b, add(yo, 0x40)))
                    let b3 := mload(add(b, add(yo, 0x60)))
                    let b4 := mload(add(b, add(yo, 0x80)))
                    mstore(add(s, yo), and(0xffffffffffffffff, xor(b0, and(not(b1), b2))))
                    mstore(add(s, add(yo, 0x20)), and(0xffffffffffffffff, xor(b1, and(not(b2), b3))))
                    mstore(add(s, add(yo, 0x40)), and(0xffffffffffffffff, xor(b2, and(not(b3), b4))))
                    mstore(add(s, add(yo, 0x60)), and(0xffffffffffffffff, xor(b3, and(not(b4), b0))))
                    mstore(add(s, add(yo, 0x80)), and(0xffffffffffffffff, xor(b4, and(not(b0), b1))))
                }
                // ---- iota ----
                mstore(add(s, 0), xor(mload(add(s, 0)), mload(add(rc, mul(rnd, 0x20)))))
            }
        }
    }

    /// @dev Absorb `input` into a fresh SHAKE256 sponge (rate 136 bytes, domain
    ///      0x1F, pad10*1) and return the state after the final permutation.
    function _shakeAbsorb(bytes memory input) internal pure returns (uint64[25] memory st) {
        uint256 len = input.length;
        uint256 off = 0;
        // full 136-byte blocks
        while (len - off >= 136) {
            for (uint256 j = 0; j < 17; j++) {
                uint64 lane = 0;
                uint256 base = off + j * 8;
                for (uint256 b = 0; b < 8; b++) {
                    lane |= uint64(uint8(input[base + b])) << (8 * b);
                }
                st[j] ^= lane;
            }
            _keccakf(st);
            off += 136;
        }
        // final (partial) block with padding
        uint256 rem = len - off;
        bytes memory blk = new bytes(136);
        for (uint256 i = 0; i < rem; i++) {
            blk[i] = input[off + i];
        }
        blk[rem] = bytes1(uint8(blk[rem]) ^ 0x1F); // domain separation
        blk[135] = bytes1(uint8(blk[135]) ^ 0x80); // final bit of pad10*1
        for (uint256 j = 0; j < 17; j++) {
            uint64 lane = 0;
            uint256 base = j * 8;
            for (uint256 b = 0; b < 8; b++) {
                lane |= uint64(uint8(blk[base + b])) << (8 * b);
            }
            st[j] ^= lane;
        }
        _keccakf(st);
    }

    /// @dev Extract the current 136-byte rate block (little-endian lanes 0..16).
    function _rateBlock(uint64[25] memory st) internal pure returns (bytes memory out) {
        out = new bytes(136);
        for (uint256 j = 0; j < 17; j++) {
            uint64 lane = st[j];
            for (uint256 b = 0; b < 8; b++) {
                out[j * 8 + b] = bytes1(uint8(lane >> (8 * b)));
            }
        }
    }

    /// @notice Squeeze `outLen` bytes of SHAKE256(input). Exposed for KAT testing.
    function shake256(bytes memory input, uint256 outLen) public pure returns (bytes memory out) {
        uint64[25] memory st = _shakeAbsorb(input);
        out = new bytes(outLen);
        uint256 pos = 0;
        while (pos < outLen) {
            bytes memory blk = _rateBlock(st);
            uint256 take = outLen - pos;
            if (take > 136) take = 136;
            for (uint256 i = 0; i < take; i++) {
                out[pos + i] = blk[i];
            }
            pos += take;
            if (pos < outLen) _keccakf(st);
        }
    }

    // ==========================================================================
    //                             HashToPoint
    // ==========================================================================

    /// @notice HashToPoint(nonce || message) -> challenge polynomial c (512 coeffs).
    /// @dev Streams SHAKE256 two bytes at a time; keeps a big-endian 16-bit sample
    ///      when it is < 5*q = 61445, reducing it mod q. Matches the reference
    ///      hash_to_point_vartime exactly.
    function hashToPoint(bytes memory nonce, bytes memory message) public pure returns (uint256[512] memory cc) {
        bytes memory input = bytes.concat(nonce, message);
        uint64[25] memory st = _shakeAbsorb(input);
        bytes memory blk = _rateBlock(st);
        uint256 bi = 0; // byte index into current rate block (0..135)
        uint256 filled = 0;
        while (filled < 512) {
            // pull two bytes, refilling / permuting when the rate block is exhausted
            if (bi == 136) {
                _keccakf(st);
                blk = _rateBlock(st);
                bi = 0;
            }
            uint8 hi = uint8(blk[bi]);
            bi++;
            if (bi == 136) {
                _keccakf(st);
                blk = _rateBlock(st);
                bi = 0;
            }
            uint8 lo = uint8(blk[bi]);
            bi++;
            uint256 w = (uint256(hi) << 8) | uint256(lo);
            if (w < 61445) {
                cc[filled] = w % Q;
                filled++;
            }
        }
    }

    // ==========================================================================
    //                        Public key / signature decode
    // ==========================================================================

    /// @notice Decode a 897-byte Falcon-512 public key into h (512 coeffs).
    /// @return h the 512 coefficients, ok false if the encoding is invalid.
    function decodePublicKey(bytes memory pk) public pure returns (uint256[512] memory h, bool ok) {
        if (pk.length != 897 || uint8(pk[0]) != PK_HEADER) {
            return (h, false);
        }
        uint256 acc = 0;
        uint256 accLen = 0;
        uint256 u = 0;
        // 896 encoded bytes = 512 * 14 bits, exact
        for (uint256 i = 1; i < 897 && u < 512; i++) {
            acc = ((acc << 8) | uint256(uint8(pk[i]))) & 0xFFFFFFFF;
            accLen += 8;
            if (accLen >= 14) {
                accLen -= 14;
                uint256 w = (acc >> accLen) & 0x3FFF;
                if (w >= Q) return (h, false);
                h[u] = w;
                u++;
            }
        }
        if (u != 512) return (h, false);
        // trailing bits must be zero
        if ((acc & ((uint256(1) << accLen) - 1)) != 0) return (h, false);
        return (h, true);
    }

    /// @notice Decode the compressed Falcon signature bytes into s2 (512 signed coeffs).
    /// @dev Mirrors the reference comp_decode: for each coefficient read a byte
    ///      giving sign + 7 low bits, then read bits until a 1 for the high part.
    ///      Rejects "-0", overflow (>2047), running off the buffer, or nonzero
    ///      trailing bits. Requires the whole buffer to be consumed.
    /// @return s2 signed coefficients (centered), ok false if invalid.
    function decodeSignature(bytes memory sig) public pure returns (int256[512] memory s2, bool ok) {
        uint256 acc = 0;
        uint256 accLen = 0;
        uint256 v = 0;
        uint256 maxLen = sig.length;
        for (uint256 u = 0; u < 512; u++) {
            if (v >= maxLen) return (s2, false);
            acc = ((acc << 8) | uint256(uint8(sig[v]))) & 0xFFFFFFFF;
            v++;
            uint256 bb = acc >> accLen;
            uint256 s = bb & 128;
            uint256 m = bb & 127;
            // read high bits until a set bit
            while (true) {
                if (accLen == 0) {
                    if (v >= maxLen) return (s2, false);
                    acc = ((acc << 8) | uint256(uint8(sig[v]))) & 0xFFFFFFFF;
                    v++;
                    accLen = 8;
                }
                accLen--;
                if (((acc >> accLen) & 1) != 0) break;
                m += 128;
                if (m > 2047) return (s2, false);
            }
            if (s != 0 && m == 0) return (s2, false); // "-0" forbidden
            s2[u] = s != 0 ? -int256(m) : int256(m);
        }
        // unused bits in the last consumed byte must be zero
        if ((acc & ((uint256(1) << accLen) - 1)) != 0) return (s2, false);
        // the whole signature buffer must be consumed exactly
        if (v != maxLen) return (s2, false);
        return (s2, true);
    }

    // ==========================================================================
    //                         Ring arithmetic + norm
    // ==========================================================================

    /// @dev b^e mod q.
    function _modpow(uint256 b, uint256 e) internal pure returns (uint256 r) {
        r = 1;
        b %= Q;
        while (e > 0) {
            if (e & 1 == 1) r = mulmod(r, b, Q);
            b = mulmod(b, b, Q);
            e >>= 1;
        }
    }

    /// @dev In-place bit-reversal permutation of a length-512 array.
    function _bitrev(uint256[512] memory a) internal pure {
        uint256 j = 0;
        for (uint256 i = 1; i < 512; i++) {
            uint256 bit = 256; // n >> 1
            while (j & bit != 0) {
                j ^= bit;
                bit >>= 1;
            }
            j ^= bit;
            if (i < j) {
                uint256 tmp = a[i];
                a[i] = a[j];
                a[j] = tmp;
            }
        }
    }

    /// @dev Iterative in-place NTT (natural order in and out) with the given root.
    function _ntt(uint256[512] memory a, uint256 root) internal pure {
        _bitrev(a);
        for (uint256 len = 2; len <= 512; len <<= 1) {
            uint256 wlen = _modpow(root, 512 / len);
            uint256 half = len >> 1;
            for (uint256 i = 0; i < 512; i += len) {
                uint256 wn = 1;
                for (uint256 k = i; k < i + half; k++) {
                    uint256 vv = mulmod(a[k + half], wn, Q);
                    uint256 u = a[k];
                    a[k] = addmod(u, vv, Q);
                    a[k + half] = addmod(u, Q - vv, Q);
                    wn = mulmod(wn, wlen, Q);
                }
            }
        }
    }

    /// @dev Negacyclic convolution t = a*b in Z_q[x]/(x^512+1) via NTT.
    ///      Pre-twist by PSI^i, length-512 NTT (root W), pointwise multiply,
    ///      inverse NTT, post-twist by PSI^-i and multiply by n^-1.
    function _ringMul(uint256[512] memory A, uint256[512] memory B) internal pure returns (uint256[512] memory r) {
        uint256[512] memory a = A;
        uint256[512] memory b = B;
        uint256 p = 1;
        for (uint256 i = 0; i < 512; i++) {
            a[i] = mulmod(a[i] % Q, p, Q);
            b[i] = mulmod(b[i] % Q, p, Q);
            p = mulmod(p, PSI, Q);
        }
        _ntt(a, W);
        _ntt(b, W);
        for (uint256 i = 0; i < 512; i++) {
            a[i] = mulmod(a[i], b[i], Q);
        }
        _ntt(a, W_INV);
        // post-twist: r[i] = a[i] * n^-1 * PSI^-i
        uint256 pinv = N_INV; // n^-1 * PSI^-0
        for (uint256 i = 0; i < 512; i++) {
            r[i] = mulmod(a[i], pinv, Q);
            pinv = mulmod(pinv, PSI_INV, Q);
        }
    }

    // ==========================================================================
    //                              Verification
    // ==========================================================================

    /// @notice Verify a Falcon-512 signature given decomposed inputs.
    /// @param pk       897-byte NIST public key (0x09 header + packed h).
    /// @param message  the signed message bytes.
    /// @param nonce    40-byte salt r.
    /// @param compSig  compressed encoding of s2 (comp_encode output).
    /// @return ok true iff the signature is a valid Falcon-512 signature.
    function verify(
        bytes memory pk,
        bytes memory message,
        bytes memory nonce,
        bytes memory compSig
    ) public pure returns (bool ok) {
        if (nonce.length != NONCE_LEN) return false;

        (uint256[512] memory h, bool okPk) = decodePublicKey(pk);
        if (!okPk) return false;

        (int256[512] memory s2, bool okSig) = decodeSignature(compSig);
        if (!okSig) return false;

        uint256[512] memory cc = hashToPoint(nonce, message);

        // s2 reduced into [0, q)
        uint256[512] memory s2m;
        for (uint256 i = 0; i < 512; i++) {
            int256 x = s2[i] % int256(Q);
            if (x < 0) x += int256(Q);
            s2m[i] = uint256(x);
        }

        // t = s2 * h  (mod q, ring)
        uint256[512] memory t = _ringMul(s2m, h);

        // s1 = c - t, centered into (-q/2, q/2]; accumulate squared norm
        uint256 norm = 0;
        for (uint256 i = 0; i < 512; i++) {
            uint256 val = addmod(cc[i], Q - t[i], Q); // (c - t) mod q, in [0,q)
            int256 s1 = val > HQ ? int256(val) - int256(Q) : int256(val);
            norm += uint256(s1 * s1);
            int256 z = s2[i];
            norm += uint256(z * z);
            if (norm > L2BOUND) return false; // early out (also guards accumulation)
        }
        return norm <= L2BOUND;
    }

    /// @notice Verify a Falcon-512 signature from a raw NIST "signed message" (sm)
    ///         blob plus the public key. This consumes the exact bytes found in
    ///         the official KAT: sm = sigLen(2, big-endian) || nonce(40) ||
    ///         message || esig, where esig = 0x29 || compSig.
    function verifySignedMessage(bytes memory pk, bytes memory sm) public pure returns (bool ok) {
        uint256 smlen = sm.length;
        if (smlen < 2 + NONCE_LEN) return false;
        uint256 sigLen = (uint256(uint8(sm[0])) << 8) | uint256(uint8(sm[1]));
        if (sigLen > smlen - 2 - NONCE_LEN) return false;
        uint256 msgLen = smlen - 2 - NONCE_LEN - sigLen;
        if (sigLen < 1) return false;

        // esig begins at 2 + NONCE_LEN + msgLen; its first byte is the header
        uint256 esigStart = 2 + NONCE_LEN + msgLen;
        if (uint8(sm[esigStart]) != SIG_HEADER) return false;

        bytes memory nonce = new bytes(NONCE_LEN);
        for (uint256 i = 0; i < NONCE_LEN; i++) nonce[i] = sm[2 + i];

        bytes memory message = new bytes(msgLen);
        for (uint256 i = 0; i < msgLen; i++) message[i] = sm[2 + NONCE_LEN + i];

        bytes memory compSig = new bytes(sigLen - 1);
        for (uint256 i = 0; i < sigLen - 1; i++) compSig[i] = sm[esigStart + 1 + i];

        return verify(pk, message, nonce, compSig);
    }

    /// @notice State-changing wrapper: verify and record the result on-chain.
    function verifyAndRecord(bytes memory pk, bytes memory sm) external returns (bool result) {
        result = verifySignedMessage(pk, sm);
        lastResult = result;
        uint256 idx = verifyCount;
        verifyCount = idx + 1;
        emit Verified(msg.sender, result, idx);
    }
}
