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

/// @title  AerePQCKeyExchange - on-chain registry + confirmed handshake for ML-KEM-768 key establishment
/// @notice A coordination and tamper-evident ledger that lets two parties establish a POST-QUANTUM
///         shared secret with ML-KEM-768 (FIPS 203, a NIST Module-Lattice KEM) and then key encrypted
///         off-chain payloads from it. The confidentiality is quantum-safe: a "harvest-now,
///         decrypt-later" adversary that records the chain cannot recover the secret, unlike an
///         ECIES / X25519 envelope a future quantum computer would open.
///
///         FLOW (heavy KEM/KDF/AEAD math run OFF-CHAIN or in the precompile; this contract stores
///         only public keys, ciphertext HASHES, a commitment, handshake state and a timeout):
///           1. A RECIPIENT registers (or rotates) its ML-KEM-768 encapsulation key `ek` via
///              registerKemKey / rotateKemKey. The contract stores keccak256(ek) and emits the full
///              key, so any counterparty can fetch it from logs and verify it against the on-chain
///              hash.
///           2. An INITIATOR, off-chain, runs ML-KEM.Encaps(ek) -> (ciphertext c, shared secret K).
///              Where the ML-KEM precompile 0x0AE6 is activated it can recompute the deterministic
///              transcript; otherwise any FIPS-203 library produces the same (c, K). The initiator
///              derives a one-way CONFIRMATION OPENING kc = keccak256(abi.encode(CONFIRM_DOMAIN,
///              chainid, this, initiator, responder, ekHash, keccak256(c), K)) and calls
///              initiateHandshake with the full c (emitted for the recipient), keccak256(c), and the
///              confirmation COMMITMENT confirmCommit = keccak256(abi.encode(kc)). K and kc are NOT
///              posted. A timeout window is armed.
///           3. The RECIPIENT, holding the decapsulation (secret) key, fetches c from the logs,
///              checks keccak256(c) == the stored kemCtHash, runs ML-KEM.Decaps to recover K, and
///              recomputes kc. It calls confirmHandshake(id, kc). The contract checks
///              keccak256(abi.encode(kc)) == confirmCommit and only then marks the handshake Confirmed
///              and records kc. Because kc is a one-way function of the secret K and was NEVER
///              published (only its hash was), a party can produce a matching kc ONLY by knowing K,
///              i.e. by actually decapsulating the ciphertext. Confirmation is therefore an on-chain,
///              transcript-bound PROOF THAT THE CONFIRMER DERIVED THE SAME SHARED SECRET.
///           4. With the handshake Confirmed, both parties key an AEAD (e.g. AES-256-GCM) from K via a
///              KDF under a DISTINCT domain and exchange encrypted payloads off-chain. The on-chain
///              Confirmed state and revealed kc are the mutually-agreed, tamper-evident anchor.
///           5. If the recipient never confirms inside the window, the handshake can be cancelled: the
///              initiator may abort at any time while pending, and ANYONE may reap it once the deadline
///              passes (timeout-cancel). A cancelled or timed-out handshake can never later confirm.
///
///         WHY REVEALING kc IS SAFE FOR CONFIDENTIALITY. kc = keccak256(...|| K) is a one-way
///         function of K, and K is a 32-byte (256-bit) high-entropy ML-KEM shared secret, so kc is
///         neither invertible nor brute-forceable back to K. The symmetric key that actually protects
///         the off-chain payload is derived as KDF(K) under a SEPARATE domain, independent of kc.
///         Publishing kc (and its commitment) therefore leaks nothing usable about K or the payload
///         key. This is the same reasoning that keeps the KEM OFF the on-chain hot path (see below).
///
///         WHY THE KEM IS NOT RUN ON-CHAIN. The ML-KEM precompile 0x0AE6 performs a DETERMINISTIC
///         Encaps(ek, m) -> (c, K). Verifying a transcript on-chain would require supplying the
///         encapsulation coins m, from which anyone could recompute K and decrypt every payload,
///         destroying the very confidentiality this rail exists to provide. So the KEM stays
///         off-chain / in the precompile off the confidential path, and the chain holds only the
///         ciphertext hash and the commit-reveal confirmation. This mirrors the sibling
///         AerePQCEncryptedIntent rail.
///
///         HONEST SCOPE AND TRUST MODEL.
///           - DEPLOYABILITY. This contract targets AERE's ML-KEM-768 precompile at 0x0AE6, which is
///             TESTNET-ONLY: it is BUILT and KAT-verified on an isolated QBFT testnet but is NOT yet
///             activated on AERE mainnet (chain 2800). Mainnet precompile activation is founder-gated.
///             So this contract is DEPLOYABLE ON TESTNET NOW and MAINNET-PENDING-ACTIVATION. Note also
///             that the contract itself never calls 0x0AE6 (the KEM is off the hot path, above), so it
///             is functionally deployable anywhere; the 0x0AE6 dependency is on the COUNTERPARTY tooling
///             that produces the ML-KEM transcript, which is why testnet is the honest deployment
///             surface today.
///           - CONSENSUS IS UNCHANGED. AERE consensus remains classical ECDSA QBFT (Foundation
///             validators). What is post-quantum here is the CONFIDENTIALITY of the established secret
///             (ML-KEM key transport). This contract adds NO consensus claim and no post-quantum
///             consensus claim.
///           - WHAT THE CHAIN GUARANTEES. (a) A recipient's advertised ML-KEM key is bound to an
///             on-chain hash any counterparty can verify; (b) a handshake references a specific
///             ciphertext (by hash) and a specific recipient key epoch, so a stale or rotated key
///             cannot be silently substituted; (c) confirmation is possible ONLY by a party that
///             derived the same shared secret (commit-reveal of a one-way opening), giving both sides a
///             non-repudiable record that the KEM completed; (d) an unconfirmed handshake always
///             resolves via timeout-cancel, so neither side is left in limbo.
///           - WHAT THE CHAIN DOES NOT DO. It does not transport, decrypt, or escrow the payload, does
///             not run the KEM, and cannot force either party to actually use K honestly off-chain. It
///             is a key-establishment COORDINATION and CONFIRMATION layer, not a message bus.
contract AerePQCKeyExchange {
    // ---- ML-KEM-768 (FIPS 203) precompile: reference constants -----------------
    /// @notice Address of AERE's ML-KEM-768 deterministic-Encaps precompile (input ek(1184)||m(32),
    ///         output c(1088)||K(32)). TESTNET-ONLY today: activated on an isolated QBFT testnet and
    ///         KAT-verified, NOT yet activated on AERE mainnet chain 2800 (founder-gated). Recorded
    ///         here for off-chain / audit reference; this contract does NOT call it (KEM stays off the
    ///         confidential hot path, see the contract NatSpec).
    address public constant ML_KEM_768_PRECOMPILE = 0x0000000000000000000000000000000000000AE6;
    /// @notice ML-KEM-768 encapsulation-key (public key) length in bytes, FIPS 203.
    uint256 public constant ML_KEM_768_EK_LEN = 1184;
    /// @notice ML-KEM-768 KEM-ciphertext length in bytes, FIPS 203.
    uint256 public constant ML_KEM_768_CT_LEN = 1088;
    /// @notice ML-KEM-768 shared-secret length in bytes, FIPS 203.
    uint256 public constant ML_KEM_768_SS_LEN = 32;

    /// @notice Domain separator for the confirmation opening kc, distinct from any other AERE PQC
    ///         domain so a value derived here can never collide with a payload-encryption key or a
    ///         signature challenge. Clients fold this (plus the transcript and K) into kc.
    bytes32 public constant CONFIRM_DOMAIN = keccak256("AerePQCKeyExchange.v1.confirm");

    // ---- Handshake timeout window ---------------------------------------------
    uint64 public constant MIN_HANDSHAKE_WINDOW = 1 minutes;
    uint64 public constant MAX_HANDSHAKE_WINDOW = 30 days;

    enum HandshakeState {
        None,      // never created
        Initiated, // ciphertext + commitment posted; awaiting the recipient's confirm
        Confirmed, // recipient revealed a matching opening; shared secret established (terminal)
        Cancelled  // aborted by the initiator, or reaped after timeout (terminal)
    }

    /// @notice A party's currently advertised ML-KEM-768 encapsulation key.
    struct KemKey {
        bytes32 keyHash; // keccak256 of the advertised ek
        uint64 epoch;    // increments on each register/rotate so clients can detect rotation
        bool active;     // false once retired; no new handshake may target it
    }

    /// @notice One handshake's on-chain state. The ciphertext bytes live in the initiate calldata and
    ///         event log, NOT in storage; storage keeps only their hash plus the commitment.
    struct Handshake {
        address initiator;     // the party that encapsulated and posted the ciphertext
        address responder;     // the recipient key owner: the only address that may confirm
        uint64 initiatedAt;    // block.timestamp at initiate
        uint64 deadline;       // confirm must land at or before this; after it, timeout-cancel is open
        HandshakeState state;
        bytes32 ekHash;        // the recipient encapsulation key epoch the initiator encapsulated to
        bytes32 kemCtHash;     // keccak256 of the ML-KEM-768 ciphertext c
        bytes32 confirmCommit; // keccak256(abi.encode(kc)); kc is the one-way, K-derived opening
        bytes32 confirmOpening;// kc, revealed by the responder at confirm (0 until Confirmed)
    }

    /// @notice party address => its advertised encapsulation key.
    mapping(address => KemKey) public kemKeyOf;
    /// @notice handshakeId => handshake record.
    mapping(uint256 => Handshake) internal _handshakes;
    /// @notice next handshake id to assign (ids start at 1; 0 is reserved for "none").
    uint256 public nextHandshakeId = 1;

    // ---- Events ----------------------------------------------------------------
    event KemKeyRegistered(address indexed party, bytes32 indexed keyHash, uint64 epoch, bytes ek);
    event KemKeyRotated(address indexed party, bytes32 indexed keyHash, uint64 epoch, bytes ek);
    event KemKeyRetired(address indexed party, bytes32 indexed keyHash);
    event HandshakeInitiated(
        uint256 indexed handshakeId,
        address indexed initiator,
        address indexed responder,
        bytes32 ekHash,
        bytes32 kemCtHash,
        bytes32 confirmCommit,
        uint64 deadline,
        bytes kemCiphertext
    );
    event HandshakeConfirmed(uint256 indexed handshakeId, address indexed responder, bytes32 confirmOpening);
    event HandshakeCancelled(uint256 indexed handshakeId, address indexed by, bool timedOut);

    // ---- Errors ----------------------------------------------------------------
    error BadEncapKeyLength(uint256 got, uint256 want);
    error KeyAlreadyActive();
    error NoActiveKey(address party);
    error ResponderKeyInactive(address responder);
    error StaleEncapKey(bytes32 got, bytes32 current);
    error BadCiphertextLength(uint256 got, uint256 want);
    error ZeroCommitment();
    error HandshakeWindowOutOfRange(uint64 got);
    error UnknownHandshake(uint256 handshakeId);
    error WrongState(HandshakeState got, HandshakeState want);
    error NotResponder(address caller, address responder);
    error HandshakeExpired(uint64 deadline, uint256 nowTs);
    error ConfirmMismatch(bytes32 got, bytes32 want);
    error NotCancellable(uint256 handshakeId);

    // ============================================================================
    //  ML-KEM key registry (register / rotate / retire)
    // ============================================================================

    /// @notice Register the caller's ML-KEM-768 encapsulation key for the FIRST time (or again after a
    ///         retire). Stores keccak256(ek) and emits the full key so counterparties can fetch and
    ///         verify it trustlessly from logs. Reverts if the caller already has an ACTIVE key; use
    ///         rotateKemKey to replace an active key in place.
    /// @param ek the FIPS 203 ML-KEM-768 encapsulation key (exactly 1184 bytes).
    function registerKemKey(bytes calldata ek) external returns (bytes32 keyHash) {
        if (kemKeyOf[msg.sender].active) revert KeyAlreadyActive();
        keyHash = _setKemKey(ek);
        emit KemKeyRegistered(msg.sender, keyHash, kemKeyOf[msg.sender].epoch, ek);
    }

    /// @notice Rotate the caller's ACTIVE ML-KEM-768 key in place, bumping its epoch. Any handshake
    ///         already Initiated against the OLD epoch keeps its own bound ekHash and can still be
    ///         confirmed (the caller must retain the matching decapsulation key), so an in-flight
    ///         handshake is never griefed by a rotation. New handshakes must target the new key.
    /// @param ek the new FIPS 203 ML-KEM-768 encapsulation key (exactly 1184 bytes).
    function rotateKemKey(bytes calldata ek) external returns (bytes32 keyHash) {
        if (!kemKeyOf[msg.sender].active) revert NoActiveKey(msg.sender);
        keyHash = _setKemKey(ek);
        emit KemKeyRotated(msg.sender, keyHash, kemKeyOf[msg.sender].epoch, ek);
    }

    /// @notice Retire the caller's advertised key so no NEW handshake can target it. Handshakes already
    ///         Initiated against it are unaffected (they carry their own ekHash) and can still be
    ///         confirmed or cancelled.
    function retireKemKey() external {
        KemKey storage k = kemKeyOf[msg.sender];
        if (!k.active) revert NoActiveKey(msg.sender);
        k.active = false;
        emit KemKeyRetired(msg.sender, k.keyHash);
    }

    /// @dev Shared register/rotate body: length-check, hash, bump epoch, mark active.
    function _setKemKey(bytes calldata ek) internal returns (bytes32 keyHash) {
        if (ek.length != ML_KEM_768_EK_LEN) revert BadEncapKeyLength(ek.length, ML_KEM_768_EK_LEN);
        keyHash = keccak256(ek);
        KemKey storage k = kemKeyOf[msg.sender];
        k.keyHash = keyHash;
        k.epoch += 1;
        k.active = true;
    }

    // ============================================================================
    //  Handshake lifecycle (initiate / confirm / cancel)
    // ============================================================================

    /// @notice Initiate a handshake to `responder` by posting an ML-KEM-768 ciphertext (encapsulated
    ///         off-chain to the responder's advertised key) plus a confirmation commitment. The full
    ///         `kemCiphertext` is passed in calldata and emitted for the responder to fetch; only its
    ///         hash and the commitment are stored. Arms a timeout window.
    /// @param responder     the recipient (must have an active advertised key).
    /// @param ekHash        keccak256 of the responder key the initiator encapsulated to; must equal
    ///                      the responder's CURRENT advertised keyHash (guards against encrypting to a
    ///                      stale/rotated key).
    /// @param kemCiphertext the ML-KEM-768 ciphertext c (exactly 1088 bytes).
    /// @param confirmCommit keccak256(abi.encode(kc)), where kc is the one-way, K-derived confirmation
    ///                      opening (see contract NatSpec). Nonzero. Never post kc or K here.
    /// @param window        seconds the responder has to confirm, in
    ///                      [MIN_HANDSHAKE_WINDOW, MAX_HANDSHAKE_WINDOW].
    function initiateHandshake(
        address responder,
        bytes32 ekHash,
        bytes calldata kemCiphertext,
        bytes32 confirmCommit,
        uint64 window
    ) external returns (uint256 handshakeId) {
        KemKey memory rk = kemKeyOf[responder];
        if (!rk.active) revert ResponderKeyInactive(responder);
        if (ekHash != rk.keyHash) revert StaleEncapKey(ekHash, rk.keyHash);
        if (kemCiphertext.length != ML_KEM_768_CT_LEN) {
            revert BadCiphertextLength(kemCiphertext.length, ML_KEM_768_CT_LEN);
        }
        if (confirmCommit == bytes32(0)) revert ZeroCommitment();
        if (window < MIN_HANDSHAKE_WINDOW || window > MAX_HANDSHAKE_WINDOW) {
            revert HandshakeWindowOutOfRange(window);
        }

        uint64 deadline = uint64(block.timestamp) + window;
        bytes32 kemCtHash = keccak256(kemCiphertext);

        handshakeId = nextHandshakeId++;
        _handshakes[handshakeId] = Handshake({
            initiator: msg.sender,
            responder: responder,
            initiatedAt: uint64(block.timestamp),
            deadline: deadline,
            state: HandshakeState.Initiated,
            ekHash: ekHash,
            kemCtHash: kemCtHash,
            confirmCommit: confirmCommit,
            confirmOpening: bytes32(0)
        });

        emit HandshakeInitiated(
            handshakeId, msg.sender, responder, ekHash, kemCtHash, confirmCommit, deadline, kemCiphertext
        );
    }

    /// @notice Confirm a handshake by revealing the confirmation opening kc. Only the designated
    ///         responder may call, and only at or before the deadline. The contract checks
    ///         keccak256(abi.encode(kc)) == the stored commitment; on match the handshake becomes
    ///         Confirmed and kc is recorded. Because kc was never published (only its hash was) and is
    ///         a one-way function of the shared secret K, a matching kc proves the caller derived the
    ///         same K, i.e. actually decapsulated the ciphertext. A wrong opening reverts and leaves
    ///         the handshake untouched (the responder may retry within the window).
    /// @param handshakeId the handshake to confirm.
    /// @param kc          the confirmation opening: keccak256(abi.encode(CONFIRM_DOMAIN, chainid, this,
    ///                    initiator, responder, ekHash, kemCtHash, K)); see contract NatSpec.
    function confirmHandshake(uint256 handshakeId, bytes32 kc) external {
        Handshake storage h = _handshakes[handshakeId];
        if (h.state == HandshakeState.None) revert UnknownHandshake(handshakeId);
        if (h.state != HandshakeState.Initiated) revert WrongState(h.state, HandshakeState.Initiated);
        if (msg.sender != h.responder) revert NotResponder(msg.sender, h.responder);
        if (block.timestamp > h.deadline) revert HandshakeExpired(h.deadline, block.timestamp);

        bytes32 got = keccak256(abi.encode(kc));
        if (got != h.confirmCommit) revert ConfirmMismatch(got, h.confirmCommit);

        h.state = HandshakeState.Confirmed;
        h.confirmOpening = kc;

        emit HandshakeConfirmed(handshakeId, msg.sender, kc);
    }

    /// @notice Cancel a pending (Initiated) handshake. Two cases:
    ///           - the INITIATOR may abort at any time while the handshake is still pending; or
    ///           - ANYONE may reap it once block.timestamp > deadline (permissionless timeout-cancel).
    ///         A Confirmed or already-Cancelled handshake cannot be cancelled. After cancel the
    ///         handshake is terminal and can never confirm, so a set of handshake parameters is
    ///         single-use.
    function cancelHandshake(uint256 handshakeId) external {
        Handshake storage h = _handshakes[handshakeId];
        if (h.state == HandshakeState.None) revert UnknownHandshake(handshakeId);
        if (h.state != HandshakeState.Initiated) revert WrongState(h.state, HandshakeState.Initiated);

        bool timedOut = block.timestamp > h.deadline;
        // Only the initiator may abort early; anyone may reap a timed-out handshake.
        if (!timedOut && msg.sender != h.initiator) revert NotCancellable(handshakeId);

        h.state = HandshakeState.Cancelled;
        emit HandshakeCancelled(handshakeId, msg.sender, timedOut);
    }

    // ============================================================================
    //  Views
    // ============================================================================

    /// @notice The full handshake record (reverts on an unknown id).
    function getHandshake(uint256 handshakeId) external view returns (Handshake memory) {
        Handshake memory h = _handshakes[handshakeId];
        if (h.state == HandshakeState.None) revert UnknownHandshake(handshakeId);
        return h;
    }

    /// @notice The state of a handshake (HandshakeState.None for an unknown id).
    function handshakeState(uint256 handshakeId) external view returns (HandshakeState) {
        return _handshakes[handshakeId].state;
    }

    /// @notice Whether a timeout-cancel would currently succeed for this handshake (Initiated and past
    ///         its deadline).
    function isTimedOut(uint256 handshakeId) external view returns (bool) {
        Handshake memory h = _handshakes[handshakeId];
        return h.state == HandshakeState.Initiated && block.timestamp > h.deadline;
    }

    /// @notice The commitment a client must post for a given confirmation opening kc, for parity checks
    ///         against off-chain code. Equals what confirmHandshake reproduces from a revealed kc.
    function confirmCommitFor(bytes32 kc) external pure returns (bytes32) {
        return keccak256(abi.encode(kc));
    }
}
