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

/**
 * @title AereMLKEM768 - on-chain ML-KEM-768 (FIPS 203) encapsulation verifier
 * @notice Thin view wrapper over the AERE native ML-KEM-768 precompile at 0x0AE6.
 *         ML-KEM-768 is AERE's post-quantum CONFIDENTIALITY primitive: every other
 *         native PQC precompile is a signature or hash (post-quantum authenticity),
 *         this one makes a Module-Lattice KEM key-agreement transcript verifiable
 *         on chain 2800.
 *
 *         The precompile performs a DETERMINISTIC ML-KEM.Encaps: given an
 *         encapsulation key `ek` (1184 bytes) and 32-byte encapsulation coins `m`,
 *         it returns the ciphertext `ct` (1088 bytes) and shared secret `ss` (32
 *         bytes) that FIPS-203 Encaps(ek, m) produces. A relying contract compares
 *         the recomputed (ct, ss) against a claimed transcript; equality proves the
 *         KEM step was performed honestly with the stated coins. Serves the UMBRA
 *         PQXDH handshake settlement path and the AERE PQC key registry.
 *
 * @dev The precompile returns EMPTY (0 bytes) for any malformed input (wrong
 *      length). Callers must treat an empty return as failure.
 *
 * @dev AUD-CRYPTO-4 - THE SHARED SECRET RETURNED HERE IS PUBLIC. Encapsulation is
 *      DETERMINISTIC in the public inputs (ek, m), both supplied as plain calldata,
 *      so the shared secret `ss` (FIPS-203 "K") is fully derivable by any observer
 *      and returning it in returndata leaks nothing beyond those public inputs. This
 *      contract is a TRANSCRIPT-VERIFICATION tool only (prove that Encaps(ek, m)
 *      yields a claimed (ct, ss)); it is NOT a live key-agreement channel. Do NOT
 *      feed it secret coins `m` or treat `ss` as confidential key material. The
 *      precompile does not perform an explicit FIPS-203 ek validity round-trip; this
 *      is harmless for the transcript-equality use, but callers must not rely on it
 *      to screen a malformed ek.
 */
contract AereMLKEM768 {
    /// @dev Native precompile address (AERE Besu PQC fork).
    address public constant MLKEM768 = 0x0000000000000000000000000000000000000AE6;

    uint256 public constant EK_LEN = 1184; // encapsulation key
    uint256 public constant M_LEN = 32; // coins
    uint256 public constant CT_LEN = 1088; // ciphertext
    uint256 public constant SS_LEN = 32; // shared secret

    /**
     * @notice Deterministically encapsulate to `ek` using coins `m`.
     * @param ek 1184-byte ML-KEM-768 encapsulation key.
     * @param m  32-byte encapsulation randomness (coins).
     * @return ok true iff the precompile returned a well-formed (ct, ss).
     * @return ct 1088-byte ciphertext.
     * @return ss 32-byte shared secret.
     */
    function encapsulate(bytes memory ek, bytes memory m)
        public
        view
        returns (bool ok, bytes memory ct, bytes memory ss)
    {
        if (ek.length != EK_LEN || m.length != M_LEN) {
            return (false, "", "");
        }
        bytes memory input = bytes.concat(ek, m);
        (bool success, bytes memory out) = MLKEM768.staticcall(input);
        if (!success || out.length != CT_LEN + SS_LEN) {
            return (false, "", "");
        }
        ct = new bytes(CT_LEN);
        ss = new bytes(SS_LEN);
        for (uint256 i = 0; i < CT_LEN; i++) {
            ct[i] = out[i];
        }
        for (uint256 i = 0; i < SS_LEN; i++) {
            ss[i] = out[CT_LEN + i];
        }
        return (true, ct, ss);
    }

    /**
     * @notice Verify a claimed ML-KEM-768 transcript (ek, m -> ct, ss).
     * @return valid true iff Encaps(ek, m) reproduces exactly the claimed ct and ss.
     */
    function verifyTranscript(
        bytes memory ek,
        bytes memory m,
        bytes memory claimedCt,
        bytes memory claimedSs
    ) external view returns (bool valid) {
        (bool ok, bytes memory ct, bytes memory ss) = encapsulate(ek, m);
        if (!ok) return false;
        return keccak256(ct) == keccak256(claimedCt) && keccak256(ss) == keccak256(claimedSs);
    }

    /**
     * @notice Recompute only the shared secret for (ek, m). Convenience for a
     *         relying contract that already holds the ciphertext off-chain.
     * @dev    AUD-CRYPTO-4 footgun: the returned `ss` is PUBLIC (deterministic in the
     *         public inputs ek, m). This is a transcript-verification convenience, NOT
     *         a confidential key-agreement primitive; never pass secret coins `m`.
     */
    function sharedSecret(bytes memory ek, bytes memory m)
        external
        view
        returns (bool ok, bytes memory ss)
    {
        bytes memory ct;
        (ok, ct, ss) = encapsulate(ek, m);
    }
}
