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

/**
 * @title AereSphincsVerifier - on-chain NIST FIPS 205 SLH-DSA (SPHINCS+) signature verification
 * @notice Verifies an SLH-DSA-SHA2-128s signature (a.k.a. SPHINCS+-SHA2-128s-simple:
 *         the smallest-signature standardized FIPS 205 parameter set) exactly per FIPS 205
 *         Algorithm 20 slh_verify_internal, using only SHA-256 (the EVM precompile at 0x02).
 *
 *         SLH-DSA is a STATELESS hash-based signature scheme. Unlike XMSS (stateful,
 *         one-time-per-leaf), a SLH-DSA key can sign an effectively unbounded number of
 *         messages with no signer state, because the leaf is chosen pseudo-randomly by
 *         a few-times FORS layer underneath a hypertree. Security rests only on the
 *         pre-image / collision resistance of SHA-256, which Shor's algorithm does not
 *         break, so it is post-quantum.
 *
 *         Parameter set SLH-DSA-SHA2-128s (FIPS 205 Table 2, security category 1):
 *           n = 16, h = 63, d = 7, h' = 9, a = 12, k = 14, lg_w = 4, w = 16,
 *           len = 35 (len1 = 32, len2 = 3), m = 30, |pk| = 32, |sig| = 7856.
 *         All tweakable hashes are SHA-256 (n <= 16 => security cat 1 uses SHA-256 only).
 *
 *         Verification path (FIPS 205 Alg. 20 slh_verify_internal, internal interface,
 *         i.e. the message M is hashed directly with no context-string domain separation):
 *           1. digest = H_msg(R, PK.seed, PK.root, M)                         [Alg. 4/11.2.1]
 *              = MGF1-SHA-256( R || PK.seed || SHA-256(R || PK.seed || PK.root || M), 30 )
 *              -> (md 21B, idx_tree 7B mod 2^54, idx_leaf 2B mod 2^9)
 *           2. PK_FORS = fors_pkFromSig(SIG_FORS, md, PK.seed, ADRS)          [Alg. 17]
 *              k=14 FORS trees, each a leaf F + an a=12 auth path (H), then T_k over 14 roots.
 *              FORS indices come from base_2^12 (big-endian, FIPS 205), NOT the round-3
 *              SPHINCS+ LSB-first message_to_indices (the one change that makes an otherwise
 *              round-3 implementation reject valid FIPS 205 signatures).
 *           3. ht_verify(PK_FORS, SIG_HT, PK.seed, idx_tree, idx_leaf, PK.root) [Alg. 12]
 *              d=7 stacked XMSS layers; each = WOTS+ pkFromSig (len=35 chains) + T_len
 *              L-compress + an h'=9 Merkle auth path (H) to that subtree root, folded up
 *              the hypertree; accept iff the top node equals PK.root.
 *
 *         The tweakable hash for the SHA-2 "simple" instances (FIPS 205 sec 11.2.1) is
 *           Th(PK.seed, ADRS, M) = Trunc_16( SHA-256( PK.seed || toByte(0,48) || ADRSc || M ) )
 *         where PK.seed (16B) is zero-padded to a full 64-byte SHA-256 block and ADRSc is the
 *         22-byte compressed address (layer[1] || tree[8] || type[1] || last-3-words[12]).
 *
 *         SHA-256 is the EVM precompile at 0x02, so this is real hash-based cryptography
 *         executed on chain, not a proof-of-a-proof and no trusted prover. Validated
 *         bit-for-bit against the OFFICIAL NIST FIPS 205 ACVP SLH-DSA-sigVer vectors (see
 *         the deployments JSON); cross-checked by a from-scratch Python verifier and by the
 *         sphincs/sphincsplus reference C (FORS index extraction corrected to FIPS 205
 *         base_2^b), both of which also match NIST 14/14.
 *
 *         HONEST SCOPE: SLH-DSA is stateless and many-time, so - unlike the live XMSS
 *         verifier - there is NO one-time-per-leaf signer-state caveat. This is a pure
 *         verifier: given (pk, message, signature) it reproduces exactly the accept/reject
 *         the FIPS 205 reference produces. Only the smallest set (SHA2-128s) is implemented;
 *         other sets (128f / 192 / 256, SHAKE) are not. No owner / admin / upgrade.
 */
contract AereSphincsVerifier {
    // ---- SLH-DSA-SHA2-128s parameters (FIPS 205 Table 2) ----
    uint256 public constant N = 16; // hash output bytes
    uint256 public constant H = 63; // total hypertree height
    uint256 public constant D = 7; // hypertree layers
    uint256 public constant HP = 9; // subtree height (h/d)
    uint256 public constant A = 12; // FORS tree height
    uint256 public constant K = 14; // FORS trees
    uint256 public constant W = 16; // Winternitz parameter
    uint256 public constant WOTS_LEN = 35; // WOTS+ chains (len1=32 + len2=3)
    uint256 public constant SIG_BYTES = 7856; // (1 + k*(a+1) + h + d*len)*n

    // signature-region byte offsets
    uint256 internal constant FORS_OFF = 16; // after R (n bytes)
    uint256 internal constant FORS_STEP = 208; // (a+1)*n per FORS tree
    uint256 internal constant HT_OFF = 2928; // 16 + k*(a+1)*n
    uint256 internal constant XMSS_STEP = 704; // (h'+len)*n per hypertree layer

    // ADRS type constants (FIPS 205 sec 4.2)
    uint256 internal constant WOTS_HASH = 0;
    uint256 internal constant WOTS_PK = 1;
    uint256 internal constant TREE = 2;
    uint256 internal constant FORS_TREE = 3;
    uint256 internal constant FORS_ROOTS = 4;

    // index masks
    uint256 internal constant TREE_MASK = (uint256(1) << 54) - 1; // 2^(h-h/d)-1
    uint256 internal constant LEAF_MASK = (uint256(1) << 9) - 1; // 2^(h/d)-1

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

    function paramSet() external pure returns (string memory) {
        return "SLH-DSA-SHA2-128s (SPHINCS+-SHA2-128s-simple), FIPS 205";
    }

    // ==========================================================================
    //                    FIPS 205 SHA-2 "simple" tweakable hash
    //   Th(PK.seed, ADRS, M) = Trunc_16(SHA-256(PK.seed || 0^48 || ADRSc || M))
    //
    //   A scratch region `sc` holds a persistent 64-byte block = PK.seed||0^48 at
    //   sc[0..64); each hash writes the 22-byte compressed ADRS at sc[64..86) then
    //   the message bytes at sc[86..), and SHA-256's the whole (86 + |M|) bytes.
    // ==========================================================================

    /// @dev pack the 22-byte compressed ADRS into the top 176 bits of a word.
    function _adc(uint256 layer, uint256 tree, uint256 typ, uint256 w5, uint256 w6, uint256 w7)
        private
        pure
        returns (uint256 adc)
    {
        // ADRSc = layer(1) || tree(8, big-endian) || type(1) || w5(4) || w6(4) || w7(4)
        adc =
            (layer << 168) |
            ((tree & 0xffffffffffffffff) << 104) |
            (typ << 96) |
            ((w5 & 0xffffffff) << 64) |
            ((w6 & 0xffffffff) << 32) |
            (w7 & 0xffffffff);
    }

    /// @dev F: single tweakable hash over one n-byte input (FORS/WOTS chain step).
    function _F(
        uint256 sc,
        uint256 layer,
        uint256 tree,
        uint256 typ,
        uint256 w5,
        uint256 w6,
        uint256 w7,
        bytes16 x
    ) private view returns (bytes16 out) {
        uint256 adc = _adc(layer, tree, typ, w5, w6, w7);
        assembly {
            let a := add(sc, 64)
            mstore(a, shl(80, adc)) // ADRSc at a[0..22), zeros a[22..32)
            mstore(add(a, 22), x) // x (16B) at a[22..38)
            // hash length = 64 (seed block) + 22 (ADRSc) + 16 (x) = 102
            if iszero(staticcall(gas(), 0x02, sc, 102, 0x00, 0x20)) { revert(0, 0) }
            out := and(mload(0x00), 0xffffffffffffffffffffffffffffffff00000000000000000000000000000000)
        }
    }

    /// @dev H: tweakable hash over two concatenated n-byte inputs (Merkle node).
    function _H(
        uint256 sc,
        uint256 layer,
        uint256 tree,
        uint256 typ,
        uint256 w5,
        uint256 w6,
        uint256 w7,
        bytes16 l,
        bytes16 r
    ) private view returns (bytes16 out) {
        uint256 adc = _adc(layer, tree, typ, w5, w6, w7);
        assembly {
            let a := add(sc, 64)
            mstore(a, shl(80, adc))
            mstore(add(a, 22), l) // l at a[22..38)
            mstore(add(a, 38), r) // r at a[38..54)
            // hash length = 64 + 22 + 32 = 118
            if iszero(staticcall(gas(), 0x02, sc, 118, 0x00, 0x20)) { revert(0, 0) }
            out := and(mload(0x00), 0xffffffffffffffffffffffffffffffff00000000000000000000000000000000)
        }
    }

    /// @dev T_l: tweakable hash over an in-memory buffer of `inLen` bytes (WOTS-pk / FORS-roots).
    function _T(
        uint256 sc,
        uint256 layer,
        uint256 tree,
        uint256 typ,
        uint256 w5,
        uint256 inPtr,
        uint256 inLen
    ) private view returns (bytes16 out) {
        uint256 adc = _adc(layer, tree, typ, w5, 0, 0);
        assembly {
            let a := add(sc, 64)
            mstore(a, shl(80, adc))
            let dst := add(a, 22)
            // portable word copy (Shanghai-safe; avoids the Cancun-only mcopy): copies
            // ceil(inLen/32) 32-byte words. inLen (224 or 560) is a multiple of 16 and the
            // source buffer is over-allocated to a 32-byte boundary, and the destination is
            // in the over-sized scratch, so a trailing partial word is harmless.
            for { let k := 0 } lt(k, inLen) { k := add(k, 32) } {
                mstore(add(dst, k), mload(add(inPtr, k)))
            }
            if iszero(staticcall(gas(), 0x02, sc, add(86, inLen), 0x00, 0x20)) { revert(0, 0) }
            out := and(mload(0x00), 0xffffffffffffffffffffffffffffffff00000000000000000000000000000000)
        }
    }

    // ==========================================================================
    //                              WOTS+ (leaf)
    // ==========================================================================

    /// @dev wots_pkFromSig (FIPS 205 Alg. 8): complete each of the 35 chains from the
    ///      signature value to chain end, then T_len-compress into the n-byte leaf.
    function _wotsPkFromSig(
        uint256 sc,
        uint256 layer,
        uint256 tree,
        uint256 idx,
        bytes calldata wsig, // 560 bytes = 35 * 16
        bytes16 mhash
    ) private view returns (bytes16) {
        // base_2^4(mhash) -> 32 message nibbles (big-endian) + checksum
        uint256[35] memory dgt;
        uint256 csum = 0;
        for (uint256 i = 0; i < 16; i++) {
            uint256 b = uint8(mhash[i]);
            uint256 hi = b >> 4;
            uint256 lo = b & 0x0f;
            dgt[2 * i] = hi;
            dgt[2 * i + 1] = lo;
            csum += (15 - hi) + (15 - lo);
        }
        csum <<= 4; // (8 - ((len2*lg_w) % 8)) % 8 == 4
        dgt[32] = (csum >> 12) & 0x0f;
        dgt[33] = (csum >> 8) & 0x0f;
        dgt[34] = (csum >> 4) & 0x0f;

        // chain-end buffer (35 * 16 = 560 bytes)
        uint256 buf;
        assembly {
            buf := mload(0x40)
            mstore(0x40, add(buf, 0x240))
        }
        for (uint256 i = 0; i < WOTS_LEN; i++) {
            bytes16 x = bytes16(wsig[i * 16:i * 16 + 16]);
            for (uint256 s = dgt[i]; s < 15; s++) {
                x = _F(sc, layer, tree, WOTS_HASH, idx, i, s, x);
            }
            assembly {
                mstore(add(buf, mul(i, 16)), x)
            }
        }
        return _T(sc, layer, tree, WOTS_PK, idx, buf, 560);
    }

    // ==========================================================================
    //                        XMSS layer (WOTS+ leaf + Merkle path)
    // ==========================================================================

    /// @dev xmss_pkFromSig (FIPS 205 Alg. 10): WOTS+ leaf then an h'=9 auth path to the
    ///      subtree root. The running tree index at level s is idx >> (s+1).
    function _xmssPkFromSig(
        uint256 sc,
        uint256 layer,
        uint256 tree,
        uint256 idx,
        bytes calldata xsig, // 704 bytes: [0:560) WOTS sig, [560:704) 9 auth nodes
        bytes16 m
    ) private view returns (bytes16 node) {
        node = _wotsPkFromSig(sc, layer, tree, idx, xsig[0:560], m);
        for (uint256 s = 0; s < HP; s++) {
            bytes16 auth = bytes16(xsig[560 + s * 16:560 + s * 16 + 16]);
            uint256 w7 = idx >> (s + 1);
            if ((idx >> s) & 1 == 0) {
                node = _H(sc, layer, tree, TREE, 0, s + 1, w7, node, auth);
            } else {
                node = _H(sc, layer, tree, TREE, 0, s + 1, w7, auth, node);
            }
        }
    }

    // ==========================================================================
    //                                FORS
    // ==========================================================================

    /// @dev fors_pkFromSig (FIPS 205 Alg. 17): reconstruct the k=14 FORS tree roots
    ///      from their leaves + auth paths, then T_k-compress into the FORS public key.
    function _forsPkFromSig(
        uint256 sc,
        uint256 idxTree,
        uint256 idxLeaf,
        uint256 mdInt, // top 168 bits hold md (base_2^12 source)
        bytes calldata sig // full signature calldata slice
    ) private view returns (bytes16) {
        uint256 buf;
        assembly {
            buf := mload(0x40)
            mstore(0x40, add(buf, 0xe0)) // 14 * 16 = 224 bytes
        }
        for (uint256 i = 0; i < K; i++) {
            uint256 fi = (mdInt >> (156 - 12 * i)) & 0xfff; // indices[i] = base_2^12
            uint256 fidx0 = (i << 12) + fi;
            uint256 base = FORS_OFF + i * FORS_STEP;
            bytes16 node = _F(sc, 0, idxTree, FORS_TREE, idxLeaf, 0, fidx0, bytes16(sig[base:base + 16]));
            for (uint256 j = 0; j < A; j++) {
                uint256 ao = base + 16 + j * 16;
                bytes16 auth = bytes16(sig[ao:ao + 16]);
                uint256 w7 = fidx0 >> (j + 1);
                if ((fi >> j) & 1 == 0) {
                    node = _H(sc, 0, idxTree, FORS_TREE, idxLeaf, j + 1, w7, node, auth);
                } else {
                    node = _H(sc, 0, idxTree, FORS_TREE, idxLeaf, j + 1, w7, auth, node);
                }
            }
            assembly {
                mstore(add(buf, mul(i, 16)), node)
            }
        }
        return _T(sc, 0, idxTree, FORS_ROOTS, idxLeaf, buf, 224);
    }

    // ==========================================================================
    //                                Verify
    // ==========================================================================

    /**
     * @notice Verify an SLH-DSA-SHA2-128s signature (FIPS 205 slh_verify_internal) as a
     *         pure/free view call.
     * @param pk         the 32-byte SLH-DSA public key (PK.seed(16) || PK.root(16))
     * @param message    the signed message bytes (internal interface: hashed directly)
     * @param signature  the 7856-byte SLH-DSA signature (R || SIG_FORS || SIG_HT)
     * @return ok        true iff the recomputed hypertree root equals PK.root
     */
    function verify(bytes calldata pk, bytes calldata message, bytes calldata signature)
        public
        view
        returns (bool ok)
    {
        if (pk.length != 32 || signature.length != SIG_BYTES) return false;

        bytes16 pkseed = bytes16(pk[0:16]);
        bytes16 pkroot = bytes16(pk[16:32]);
        bytes16 R = bytes16(signature[0:16]);

        // scratch: sc[0..64) = PK.seed || 0^48 (persistent across all tweakable hashes)
        uint256 sc;
        assembly {
            sc := mload(0x40)
            mstore(0x40, add(sc, 0x340)) // 832 bytes (86 + up-to-560 input, padded)
            mstore(sc, pkseed) // PK.seed at sc[0..16), zeros sc[16..32)
            mstore(add(sc, 16), 0)
            mstore(add(sc, 48), 0) // ensure sc[16..64) = 0
        }

        // digest = H_msg(R, PK.seed, PK.root, M)
        //        = MGF1-SHA256( R||PK.seed||SHA256(R||PK.seed||PK.root||M), 30 )
        //        = SHA256( R || PK.seed || inner || 0x00000000 )[0:30]   (30 <= 32)
        bytes32 inner = sha256(abi.encodePacked(R, pkseed, pkroot, message));
        bytes32 digest = sha256(abi.encodePacked(R, pkseed, inner, bytes4(0)));

        uint256 dig = uint256(digest);
        uint256 mdInt = dig >> 88; // top 168 bits = md (21 bytes)
        uint256 idxTree = (dig >> 32) & ((uint256(1) << 56) - 1) & TREE_MASK;
        uint256 idxLeaf = (dig >> 16) & LEAF_MASK;

        // FORS public key
        bytes16 pkFors = _forsPkFromSig(sc, idxTree, idxLeaf, mdInt, signature);

        // hypertree verification
        bytes16 node = _xmssPkFromSig(sc, 0, idxTree, idxLeaf, signature[HT_OFF:HT_OFF + XMSS_STEP], pkFors);
        uint256 tree = idxTree;
        for (uint256 layer = 1; layer < D; layer++) {
            uint256 leaf = tree & LEAF_MASK;
            tree >>= HP;
            uint256 o = HT_OFF + layer * XMSS_STEP;
            node = _xmssPkFromSig(sc, layer, tree, leaf, signature[o:o + XMSS_STEP], node);
        }
        return node == pkroot;
    }

    /**
     * @notice State-changing wrapper: verifies and records the result on-chain (event +
     *         counter). SLH-DSA-SHA2-128s verification is hash-heavy (FORS + a d=7 hypertree
     *         of WOTS+ chains); whether a full verifyAndRecord lands under the Fusaka
     *         EIP-7825 2^24 (16,777,216) per-tx gas cap is measured at deploy time and
     *         recorded in the deployments JSON. `verify()` is the free pure view and always
     *         works via eth_call regardless of the tx cap.
     */
    function verifyAndRecord(bytes calldata pk, bytes calldata message, bytes calldata signature)
        external
        returns (bool ok)
    {
        ok = verify(pk, message, signature);
        lastResult = ok;
        verifyCount += 1;
        emit Verified(msg.sender, ok, verifyCount);
    }
}
