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

/**
 * @title AerePQCPrecompiles - clean view wrappers for AERE's two newest PQC precompiles
 *
 * @notice A dependency-free Solidity library that any contract can `using`/import to call the
 *         two native precompiles added to the AERE Besu PQC fork:
 *
 *           - ML-KEM-768 (FIPS 203) deterministic encapsulation, at 0x0AE6.
 *           - Falcon HashToPoint (SHAKE256 rejection sampler), at 0x0AE7.
 *
 *         Every function is `internal view` and fail-closed: it validates the exact wire
 *         lengths, performs one `staticcall`, checks the returned length, and returns
 *         `ok == false` (never reverts) on any malformed input or unexpected output. Because
 *         the functions are `internal`, the library is inlined into the caller - there is no
 *         separate deployment and no external call surface to trust.
 *
 * @dev    HONEST SCOPE - TESTNET ONLY (as of 2026-07-18). These two precompiles are live on the
 *         AERE TESTNET ONLY; they are NOT activated on mainnet chain 2800, which has exactly the
 *         FIVE precompiles 0x0AE1..0x0AE5 (activated block 9,189,161). On a chain WITHOUT these
 *         precompiles a staticcall to 0x0AE6 / 0x0AE7 hits an empty account and returns success
 *         with EMPTY data, so the length check fails and every wrapper here returns `ok == false`.
 *         A contract using this library therefore behaves safely (deny, never a false positive)
 *         on mainnet today, and starts succeeding automatically once the precompiles are
 *         activated per docs/PQC-PRECOMPILES-MAINNET-ACTIVATION-PLAN-2026-07-18.md (founder +
 *         external-audit gated). Nothing here changes AERE consensus, which remains classical
 *         secp256k1 ECDSA QBFT.
 *
 * @dev    These wrappers deliberately do NOT hold state and make NO cryptographic claim of their
 *         own: they are a thin, exact restatement of each precompile's byte contract. The
 *         precompiles themselves are KAT/ACVP-validated (ML-KEM-768: 25/25 NIST ACVP; Falcon
 *         HashToPoint: 12/12 vs an independent FIPS 202 SHAKE256 oracle). A standalone contract
 *         form of the ML-KEM-768 wrapper also exists at AereMLKEM768.sol; this library is the
 *         reusable, importable surface covering BOTH precompiles.
 */
library AerePQCPrecompiles {
    // ---- Precompile addresses (TESTNET ONLY; not on mainnet 2800 yet) ----------------------

    /// @dev ML-KEM-768 (FIPS 203) deterministic-Encaps precompile.
    address internal constant MLKEM768 = address(0x0AE6);
    /// @dev Falcon HashToPoint (SHAKE256 rejection sampler) precompile.
    address internal constant FALCON_HASHTOPOINT = address(0x0AE7);

    // ---- ML-KEM-768 field sizes (FIPS 203, bytes) ------------------------------------------

    /// @dev Encapsulation key (public key) length.
    uint256 internal constant MLKEM768_EK_LEN = 1184;
    /// @dev Encapsulation randomness ("coins") length.
    uint256 internal constant MLKEM768_M_LEN = 32;
    /// @dev Ciphertext length.
    uint256 internal constant MLKEM768_CT_LEN = 1088;
    /// @dev Shared-secret length.
    uint256 internal constant MLKEM768_SS_LEN = 32;

    // ---- Falcon HashToPoint sizes ----------------------------------------------------------

    /// @dev Falcon salt / nonce length (bytes).
    uint256 internal constant FALCON_NONCE_LEN = 40;
    /// @dev Minimum Falcon logn accepted (n = 2^logn). Falcon-512 uses 9, Falcon-1024 uses 10.
    uint8 internal constant FALCON_MIN_LOGN = 9;
    /// @dev Maximum Falcon logn accepted.
    uint8 internal constant FALCON_MAX_LOGN = 10;
    /// @dev Falcon ring modulus q; every returned coefficient is in [0, q).
    uint256 internal constant FALCON_Q = 12289;

    // ========================================================================================
    //                                    ML-KEM-768 @ 0x0AE6
    // ========================================================================================

    /**
     * @notice Deterministically encapsulate to `ek` using the 32-byte coins `m`, via the native
     *         ML-KEM-768 precompile. Reproduces FIPS 203 `ML-KEM.Encaps(ek, m)`: the same
     *         (ct, ss) every node computes for the same (ek, m).
     *
     * @dev    Wire contract of 0x0AE6:
     *           input  = ek(1184) || m(32)   (1216 bytes, exact)
     *           output = ct(1088) || ss(32)  (1120 bytes); EMPTY on malformed input.
     *         Fail-closed: returns (false, "", "") on any wrong-length input or output (which
     *         includes a chain where the precompile is not activated).
     *
     * @param ek 1184-byte ML-KEM-768 encapsulation key (public key).
     * @param m  32-byte encapsulation randomness (coins).
     * @return ok true iff the precompile returned a well-formed 1120-byte (ct, ss).
     * @return ct 1088-byte ciphertext (empty when ok is false).
     * @return ss 32-byte shared secret (empty when ok is false).
     */
    function mlkem768Encapsulate(bytes memory ek, bytes memory m)
        internal
        view
        returns (bool ok, bytes memory ct, bytes memory ss)
    {
        if (ek.length != MLKEM768_EK_LEN || m.length != MLKEM768_M_LEN) {
            return (false, "", "");
        }
        bytes memory input = bytes.concat(ek, m);
        (bool success, bytes memory out) = MLKEM768.staticcall(input);
        if (!success || out.length != MLKEM768_CT_LEN + MLKEM768_SS_LEN) {
            return (false, "", "");
        }
        ct = new bytes(MLKEM768_CT_LEN);
        ss = new bytes(MLKEM768_SS_LEN);
        for (uint256 i = 0; i < MLKEM768_CT_LEN; i++) {
            ct[i] = out[i];
        }
        for (uint256 i = 0; i < MLKEM768_SS_LEN; i++) {
            ss[i] = out[MLKEM768_CT_LEN + i];
        }
        return (true, ct, ss);
    }

    /**
     * @notice Recompute only the 32-byte shared secret for (ek, m). Convenience for a relying
     *         contract that already holds the ciphertext.
     * @return ok true iff the precompile returned a well-formed transcript.
     * @return ss 32-byte shared secret (empty when ok is false).
     */
    function mlkem768SharedSecret(bytes memory ek, bytes memory m)
        internal
        view
        returns (bool ok, bytes memory ss)
    {
        bytes memory ct;
        (ok, ct, ss) = mlkem768Encapsulate(ek, m);
    }

    /**
     * @notice Verify a claimed ML-KEM-768 transcript: does Encaps(ek, m) reproduce exactly the
     *         claimed ciphertext and shared secret?
     * @return valid true iff the recomputed (ct, ss) equal the claimed values byte-for-byte.
     */
    function mlkem768VerifyTranscript(
        bytes memory ek,
        bytes memory m,
        bytes memory claimedCt,
        bytes memory claimedSs
    ) internal view returns (bool valid) {
        (bool ok, bytes memory ct, bytes memory ss) = mlkem768Encapsulate(ek, m);
        if (!ok) return false;
        return keccak256(ct) == keccak256(claimedCt) && keccak256(ss) == keccak256(claimedSs);
    }

    // ========================================================================================
    //                                Falcon HashToPoint @ 0x0AE7
    // ========================================================================================

    /**
     * @notice HashToPoint(nonce || message) -> the Falcon challenge polynomial c, via the native
     *         precompile. Matches the reference `hash_to_point_vartime`: absorb `nonce||message`
     *         into SHAKE256, squeeze 16-bit big-endian samples, keep `w mod q` whenever
     *         `w < 5q = 61445`, until `n = 2^logn` coefficients are collected.
     *
     * @dev    Wire contract of 0x0AE7:
     *           input  = logn(1) || nonce(40) || message(rest)
     *           output = 2n bytes = n big-endian uint16 coefficients in [0, q), q=12289; EMPTY on malformed.
     *         The raw 2n-byte output is returned as-is; use {challengeCoeff} to read a coefficient,
     *         or decode each pair as `(challenge[2i] << 8) | challenge[2i+1]`. Fail-closed: returns
     *         (false, "") on a bad logn, a wrong-length nonce, or an unexpected output length
     *         (which includes a chain where the precompile is not activated).
     *
     * @param logn    Falcon log-degree: 9 for Falcon-512 (n=512) or 10 for Falcon-1024 (n=1024).
     * @param nonce   40-byte Falcon salt r.
     * @param message the message bytes hashed together with the nonce.
     * @return ok        true iff the precompile returned the expected 2n bytes.
     * @return challenge the 2n-byte challenge polynomial (n big-endian uint16 coeffs); empty when ok is false.
     */
    function falconHashToPoint(uint8 logn, bytes memory nonce, bytes memory message)
        internal
        view
        returns (bool ok, bytes memory challenge)
    {
        if (logn < FALCON_MIN_LOGN || logn > FALCON_MAX_LOGN) return (false, "");
        if (nonce.length != FALCON_NONCE_LEN) return (false, "");
        uint256 n = uint256(1) << logn;
        uint256 outLen = 2 * n;
        bytes memory input = bytes.concat(bytes1(logn), nonce, message);
        (bool success, bytes memory out) = FALCON_HASHTOPOINT.staticcall(input);
        if (!success || out.length != outLen) return (false, "");
        return (true, out);
    }

    /**
     * @notice Read coefficient `i` of a HashToPoint challenge polynomial produced by
     *         {falconHashToPoint}. Coefficient `i` is the big-endian uint16 at bytes [2i, 2i+1].
     * @dev    Reverts if `i` is out of range for the given challenge buffer.
     * @param challenge the 2n-byte challenge polynomial.
     * @param i         coefficient index (0-based).
     * @return coeff the i-th coefficient, in [0, q).
     */
    function challengeCoeff(bytes memory challenge, uint256 i) internal pure returns (uint256 coeff) {
        uint256 hi = 2 * i;
        require(hi + 1 < challenge.length, "AerePQCPrecompiles: coeff index OOB");
        return (uint256(uint8(challenge[hi])) << 8) | uint256(uint8(challenge[hi + 1]));
    }
}
