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

/**
 * @title AereXmssVerifier - on-chain RFC 8391 XMSS many-time hash-based signature verification
 * @notice Verifies an XMSS signature for parameter set XMSS-SHA2_10_256 (OID 0x00000001:
 *         n = 32, Winternitz w = 16, tree height h = 10, hash = SHA-256) exactly per
 *         RFC 8391 / the XMSS reference implementation (github.com/XMSS/xmss-reference).
 *
 *         XMSS is the honest many-time extension of the one-time WOTS+ scheme already live
 *         on this chain (AerePQCVerifier). A signature proves a WOTS+ one-time leaf plus a
 *         Merkle authentication path from that leaf up to a long-term public root. Security
 *         rests only on the pre-image / collision resistance of SHA-256, which Shor's
 *         algorithm does not break, so it is post-quantum.
 *
 *         Verification path (RFC 8391 sections 4.1.9 / 4.1.10, single tree d = 1):
 *           1. M' = H_msg(R || root || toByte(idx, 32) || M)
 *           2. pk_ots = WOTS_PKFromSig(sig_ots, M', SEED, OTS-address(idx))   [67 chains]
 *           3. leaf   = ltree(pk_ots, SEED, LTree-address(idx))               [L-tree compress]
 *           4. node   = compute_root(leaf, idx, auth, SEED, HashTree-address) [h=10 Merkle path]
 *           5. accept iff node == root
 *
 *         All hashing uses the RFC 8391 keyed, address-domain-separated toolbox with
 *         padding_len = n = 32:
 *           F(KEY, M)      = SHA256( toByte(0,32) || KEY || M )
 *           H(KEY, M)      = SHA256( toByte(1,32) || KEY || M )
 *           H_msg(KEY, M)  = SHA256( toByte(2,32) || KEY || M )
 *           PRF(KEY, ADRS) = SHA256( toByte(3,32) || KEY || ADRS )
 *         with chain/tree bitmasks and keys derived by PRF(SEED, ADRS) over the 32-byte
 *         hash address (layer/tree/type + OTS/L-tree/hash-tree fields), key_and_mask in {0,1,2}.
 *
 *         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 against
 *         the OFFICIAL deterministic reference known-answer vector (see deployments JSON).
 *
 *         HONEST SCOPE: XMSS is a STATEFUL many-time scheme. This contract verifies that a
 *         given signature is valid for a given (root, SEED); it does NOT and cannot enforce
 *         the signer's one-time-per-leaf state. Signing SECURITY requires the off-chain
 *         signer to never reuse a leaf index (idx). Reusing an index breaks WOTS+, not this
 *         verifier. This is a verifier, not a state manager.
 */
contract AereXmssVerifier {
    // ----- XMSS-SHA2_10_256 parameters (RFC 8391 OID 0x00000001) -----
    uint256 public constant N = 32; // hash output bytes
    uint256 public constant W = 16; // Winternitz parameter
    uint256 public constant LOG_W = 4; // log2(w)
    uint256 public constant LEN1 = 64; // message base-w digits (8*n / log_w)
    uint256 public constant LEN2 = 3; // checksum base-w digits
    uint256 public constant LEN = 67; // total WOTS+ chains
    uint256 public constant H = 10; // Merkle tree height
    uint32 public constant XMSS_OID = 0x00000001; // XMSS-SHA2_10_256

    // hash address types (RFC 8391 section 2.5)
    uint256 internal constant ADDR_OTS = 0;
    uint256 internal constant ADDR_LTREE = 1;
    uint256 internal constant ADDR_HASHTREE = 2;

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

    // ==========================================================================
    //                    RFC 8391 keyed hash toolbox (SHA-256)
    // ==========================================================================

    /// @dev Pack an XMSS 32-byte hash address. layer (word0) and tree (words 1-2) are
    ///      always zero for the single-tree XMSS-SHA2_10_256 set; type is word3, then
    ///      the type-specific words 4/5/6 and key_and_mask in word7.
    function _addr(uint256 typ, uint256 a4, uint256 a5, uint256 a6, uint256 km)
        internal
        pure
        returns (bytes32)
    {
        return bytes32((typ << 128) | (a4 << 96) | (a5 << 64) | (a6 << 32) | km);
    }

    /// @dev PRF(SEED, ADRS) = SHA256( toByte(3,32) || SEED || ADRS ).
    function _prf(bytes32 seed, bytes32 addr) internal pure returns (bytes32) {
        return sha256(abi.encodePacked(bytes32(uint256(3)), seed, addr));
    }

    /// @dev RFC 8391 F: one hash-chain step. Derives a per-address key and bitmask by PRF,
    ///      then F(KEY, in ^ mask) = SHA256( toByte(0,32) || KEY || (in ^ mask) ).
    function _f(bytes32 seed, uint256 typ, uint256 a4, uint256 a5, uint256 a6, bytes32 x)
        internal
        pure
        returns (bytes32)
    {
        bytes32 key = _prf(seed, _addr(typ, a4, a5, a6, 0));
        bytes32 mask = _prf(seed, _addr(typ, a4, a5, a6, 1));
        return sha256(abi.encodePacked(bytes32(0), key, x ^ mask));
    }

    /// @dev RFC 8391 RAND_HASH / H over two n-byte children. Key + two n-byte masks by PRF,
    ///      then H(KEY, (L ^ m0) || (R ^ m1)) = SHA256( toByte(1,32) || KEY || (L^m0) || (R^m1) ).
    function _h(
        bytes32 seed,
        uint256 typ,
        uint256 a4,
        uint256 a5,
        uint256 a6,
        bytes32 left,
        bytes32 right
    ) internal pure returns (bytes32) {
        bytes32 key = _prf(seed, _addr(typ, a4, a5, a6, 0));
        bytes32 m0 = _prf(seed, _addr(typ, a4, a5, a6, 1));
        bytes32 m1 = _prf(seed, _addr(typ, a4, a5, a6, 2));
        return sha256(abi.encodePacked(bytes32(uint256(1)), key, left ^ m0, right ^ m1));
    }

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

    /// @notice Derive the 67 base-16 chain lengths (64 message digits + 3 checksum digits)
    ///         from a 32-byte message hash, exactly per RFC 8391 chain_lengths.
    function chainLengths(bytes32 m) public pure returns (uint256[67] memory d) {
        uint256 csum = 0;
        for (uint256 i = 0; i < 32; i++) {
            uint8 b = uint8(m[i]);
            uint256 hi = b >> 4;
            uint256 lo = b & 0x0f;
            d[2 * i] = hi;
            d[2 * i + 1] = lo;
            csum += (15 - hi) + (15 - lo);
        }
        // csum << (8 - (len2*log_w % 8)) == csum << 4, then base_w over 2 big-endian bytes.
        uint256 c = csum << 4;
        d[64] = (c >> 12) & 0x0f;
        d[65] = (c >> 8) & 0x0f;
        d[66] = (c >> 4) & 0x0f;
    }

    /// @dev WOTS_PKFromSig: complete each of the 67 chains from the signature value to the
    ///      chain end (position w-1), returning the 67 WOTS+ public-key chain values.
    function _wotsPkFromSig(bytes32 seed, uint256 idxLeaf, bytes32 mhash, bytes32[67] calldata sig)
        internal
        pure
        returns (bytes32[67] memory pk)
    {
        uint256[67] memory lens = chainLengths(mhash);
        for (uint256 i = 0; i < LEN; i++) {
            bytes32 x = sig[i];
            // gen_chain: for s in [lens[i], w-1): F with hash-address s, chain-address i.
            for (uint256 s = lens[i]; s < W - 1; s++) {
                x = _f(seed, ADDR_OTS, idxLeaf, i, s, x);
            }
            pk[i] = x;
        }
    }

    /// @dev L-tree: compress the 67 WOTS+ public-key values into a single n-byte leaf.
    function _lTree(bytes32 seed, uint256 idxLeaf, bytes32[67] memory pk)
        internal
        pure
        returns (bytes32)
    {
        // Work on a mutable copy sized to LEN; unused tail is ignored via the shrinking `l`.
        bytes32[67] memory nodes = pk;
        uint256 l = LEN;
        uint256 height = 0;
        while (l > 1) {
            uint256 parent = l >> 1;
            for (uint256 i = 0; i < parent; i++) {
                nodes[i] = _h(seed, ADDR_LTREE, idxLeaf, height, i, nodes[2 * i], nodes[2 * i + 1]);
            }
            if (l & 1 == 1) {
                nodes[parent] = nodes[l - 1];
                l = parent + 1;
            } else {
                l = parent;
            }
            height++;
        }
        return nodes[0];
    }

    // ==========================================================================
    //                        Merkle authentication path
    // ==========================================================================

    /// @dev compute_root: fold the leaf with the h=10 authentication path up to the root,
    ///      using the hash-tree address (type 2) with the correct per-level height/index.
    function _computeRoot(bytes32 seed, bytes32 leaf, uint256 leafIdx, bytes32[10] calldata auth)
        internal
        pure
        returns (bytes32)
    {
        bytes32 left;
        bytes32 right;
        if (leafIdx & 1 == 1) {
            left = auth[0];
            right = leaf;
        } else {
            left = leaf;
            right = auth[0];
        }
        uint256 li = leafIdx;
        for (uint256 i = 0; i < H - 1; i++) {
            li >>= 1;
            bytes32 node = _h(seed, ADDR_HASHTREE, 0, i, li, left, right);
            if (li & 1 == 1) {
                left = auth[i + 1];
                right = node;
            } else {
                left = node;
                right = auth[i + 1];
            }
        }
        li >>= 1;
        return _h(seed, ADDR_HASHTREE, 0, H - 1, li, left, right);
    }

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

    /**
     * @notice Verify an XMSS-SHA2_10_256 signature (RFC 8391) as a pure/free view call.
     * @param pubRoot  the long-term XMSS public root (first 32 bytes of the XMSS public key)
     * @param pubSeed  the XMSS public SEED (last 32 bytes of the XMSS public key)
     * @param idx      the leaf index used by the signer (0 .. 2^h - 1)
     * @param R        the per-signature randomizer (from the signature)
     * @param message  the signed message bytes
     * @param wotsSig  the 67 WOTS+ one-time-signature chain values
     * @param auth     the h = 10 Merkle authentication-path nodes
     * @return ok      true iff the recomputed root equals pubRoot
     */
    function verify(
        bytes32 pubRoot,
        bytes32 pubSeed,
        uint32 idx,
        bytes32 R,
        bytes calldata message,
        bytes32[67] calldata wotsSig,
        bytes32[10] calldata auth
    ) public pure returns (bool ok) {
        // M' = H_msg( toByte(2,32) || R || root || toByte(idx, 32) || M )
        bytes32 mhash = sha256(
            abi.encodePacked(bytes32(uint256(2)), R, pubRoot, bytes32(uint256(idx)), message)
        );
        uint256 idxLeaf = uint256(idx) & ((1 << H) - 1);
        bytes32[67] memory wpk = _wotsPkFromSig(pubSeed, idxLeaf, mhash, wotsSig);
        bytes32 leaf = _lTree(pubSeed, idxLeaf, wpk);
        bytes32 root = _computeRoot(pubSeed, leaf, idxLeaf, auth);
        return root == pubRoot;
    }

    /**
     * @notice State-changing wrapper: verifies and records the result on-chain (event + counter).
     *         XMSS verification is SHA-256-heavy but far cheaper than a lattice verify, so this
     *         lands under the Fusaka EIP-7825 2^24 (16,777,216) per-transaction gas cap and runs
     *         as an ordinary transaction. `verify()` remains the free view for pre-flight.
     */
    function verifyAndRecord(
        bytes32 pubRoot,
        bytes32 pubSeed,
        uint32 idx,
        bytes32 R,
        bytes calldata message,
        bytes32[67] calldata wotsSig,
        bytes32[10] calldata auth
    ) external returns (bool ok) {
        ok = verify(pubRoot, pubSeed, idx, R, message, wotsSig, auth);
        lastResult = ok;
        verifyCount += 1;
        emit Verified(msg.sender, pubRoot, ok, verifyCount);
    }
}
