// SPDX-License-Identifier: LGPL-3.0-only
pragma solidity ^0.8.23;

/// @title  AerePQCThreshold
/// @notice Shared post-quantum verification + envelope construction, extracted VERBATIM from the
///         proven `AereThresholdPQCRegistry` (17/17 tests, precompile envelope checked
///         byte-for-byte against the LIVE mainnet AERE precompiles 0x0AE1..0x0AE4). Both the
///         registry-style flows and the ERC-4337 `AereThresholdAccount` route through this ONE
///         verification path, so there is a single audited surface (avoids re-implementing the
///         crypto envelope and re-introducing a domain-confusion / parsing bug).
/// @dev    NIST PQC schemes: 1=Falcon-512 (0x0AE1), 2=Falcon-1024 (0x0AE2), 3=ML-DSA-44 (0x0AE3,
///         FIPS 204), 4=SLH-DSA-128s (0x0AE4, FIPS 205). Application/account layer only; AERE
///         consensus is classical ECDSA QBFT.
library AerePQCThreshold {
    uint8 internal constant SCHEME_FALCON512 = 1;
    uint8 internal constant SCHEME_FALCON1024 = 2;
    uint8 internal constant SCHEME_MLDSA44 = 3;
    uint8 internal constant SCHEME_SLHDSA128S = 4;

    address internal constant PRECOMPILE_FALCON512 = address(0x0AE1);
    address internal constant PRECOMPILE_FALCON1024 = address(0x0AE2);
    address internal constant PRECOMPILE_MLDSA44 = address(0x0AE3);
    address internal constant PRECOMPILE_SLHDSA128S = address(0x0AE4);

    uint256 internal constant FALCON512_PK_LEN = 897;
    uint256 internal constant FALCON1024_PK_LEN = 1793;
    uint256 internal constant MLDSA44_PK_LEN = 1312;
    uint256 internal constant SLHDSA128S_PK_LEN = 32;
    uint8 internal constant FALCON512_PK_HEADER = 0x09; // 0x00 | logn=9
    uint8 internal constant FALCON1024_PK_HEADER = 0x0A; // 0x00 | logn=10

    uint256 internal constant MLDSA44_SIG_LEN = 2420;
    uint256 internal constant SLHDSA128S_SIG_LEN = 7856;
    uint256 internal constant FALCON_NONCE_LEN = 40;

    /// @notice True iff `pk` is a structurally valid public key for `scheme` (length + Falcon
    ///         header). A malformed key must never be accepted into a committee.
    function validPubKey(uint8 scheme, bytes memory pk) internal pure returns (bool) {
        if (scheme == SCHEME_FALCON512) return pk.length == FALCON512_PK_LEN && uint8(pk[0]) == FALCON512_PK_HEADER;
        if (scheme == SCHEME_FALCON1024) return pk.length == FALCON1024_PK_LEN && uint8(pk[0]) == FALCON1024_PK_HEADER;
        if (scheme == SCHEME_MLDSA44) return pk.length == MLDSA44_PK_LEN;
        if (scheme == SCHEME_SLHDSA128S) return pk.length == SLHDSA128S_PK_LEN;
        return false;
    }

    /// @dev Assemble (precompile, input) for a scheme, matching the live precompile spec offsets
    ///      exactly. `wellFormed` is false for a wrong-length signature (an invalid leg, never a
    ///      revert), so a malformed leg is counted as "did not verify" rather than aborting.
    function buildInput(uint8 scheme, bytes memory pubKey, bytes32 message, bytes memory signature)
        internal
        pure
        returns (address precompile, bytes memory input, bool wellFormed)
    {
        if (scheme == SCHEME_FALCON512 || scheme == SCHEME_FALCON1024) {
            if (signature.length <= FALCON_NONCE_LEN) return (address(0), "", false);
            uint256 sigLen = signature.length - FALCON_NONCE_LEN; // esig length
            if (sigLen > type(uint16).max) return (address(0), "", false);

            bytes memory nonce = slice(signature, 0, FALCON_NONCE_LEN);
            bytes memory esig = slice(signature, FALCON_NONCE_LEN, sigLen);

            // sm = sigLen(2, big-endian) || nonce(40) || message(32) || esig
            bytes memory sm = abi.encodePacked(uint16(sigLen), nonce, message, esig);
            input = abi.encodePacked(pubKey, sm);
            precompile = scheme == SCHEME_FALCON512 ? PRECOMPILE_FALCON512 : PRECOMPILE_FALCON1024;
            wellFormed = true;
        } else if (scheme == SCHEME_MLDSA44) {
            if (signature.length != MLDSA44_SIG_LEN) return (address(0), "", false);
            input = abi.encodePacked(pubKey, signature, message); // pk(1312) || sig(2420) || message(32)
            precompile = PRECOMPILE_MLDSA44;
            wellFormed = true;
        } else if (scheme == SCHEME_SLHDSA128S) {
            if (signature.length != SLHDSA128S_SIG_LEN) return (address(0), "", false);
            input = abi.encodePacked(pubKey, signature, message); // pk(32) || sig(7856) || message(32)
            precompile = PRECOMPILE_SLHDSA128S;
            wellFormed = true;
        } else {
            return (address(0), "", false);
        }
    }

    /// @notice Verify a single PQC signature over `message` in full, on-chain, via the scheme's
    ///         live precompile. Returns false (never reverts) on a malformed or rejected leg.
    function verify(uint8 scheme, bytes memory pubKey, bytes32 message, bytes memory signature)
        internal
        view
        returns (bool)
    {
        (address precompile, bytes memory input, bool wellFormed) = buildInput(scheme, pubKey, message, signature);
        if (!wellFormed) return false;
        (bool ok, bytes memory ret) = precompile.staticcall(input);
        return ok && ret.length >= 32 && ret[31] == 0x01;
    }

    function slice(bytes memory data, uint256 start, uint256 len) internal pure returns (bytes memory out) {
        out = new bytes(len);
        for (uint256 i = 0; i < len; i++) {
            out[i] = data[start + i];
        }
    }
}