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

/**
 * @title AerePQCKeyRegistry, permissionless post-quantum public-key registry with
 *        on-chain proof-of-possession (AERE chain 2800)
 *
 * @notice The anchor a wallet, agent, or app uses to PUBLISH its post-quantum public
 *         keys and prove, on-chain, that it actually holds the matching private key.
 *         Anyone can bind one or more NIST PQC public keys (Falcon-512, Falcon-1024,
 *         ML-DSA-44, SLH-DSA-128s) to their identity (an address). Registration is
 *         gated by a PROOF-OF-POSSESSION: the registrant must supply a valid PQC
 *         signature, verified in full on-chain by AERE's live native precompile, over
 *         a challenge this contract binds to (chain, contract, owner, scheme, key,
 *         per-identity nonce). A party that does not hold the private key cannot
 *         produce that signature, so a key can never be registered under an identity
 *         that does not control it, and an observed proof cannot be replayed to claim
 *         the key under a different address (the owner is inside the challenge).
 *
 *         Keys can be ROTATED (register a successor, proving possession of the NEW key,
 *         which retires the old one) and REVOKED (terminal), by the identity owner only.
 *
 *         HONEST SCOPE. Application-layer PQC verification. This does NOT change AERE
 *         consensus: blocks are still produced and signed by Besu QBFT validators with
 *         classical ECDSA (secp256k1). What is new and live is the set of native
 *         precompiles this contract calls (mainnet 2800, activation block 9,189,161),
 *         each wrapping the audited Bouncy Castle NIST verifiers. No funds are ever held.
 *
 * @dev    The precompile wire format is IDENTICAL to AerePQCAttestation (proven against
 *         the live mainnet precompiles via eth_call). See that contract's NatSpec and
 *         scripts/pqc for the byte layout; the encoding is reproduced verbatim below so
 *         a consumer can verify one file. Signature envelope handed to this contract:
 *           Falcon (1,2):     signature = nonce(40) || esig     (esig starts 0x29 / 0x2A)
 *           ML-DSA-44 (3):    signature = sig, exactly 2420 bytes
 *           SLH-DSA-128s (4): signature = sig, exactly 7856 bytes
 *         This contract supplies the 32-byte challenge/message itself, so a signer can
 *         never bind a proof to anything other than what this contract commits to.
 */
contract AerePQCKeyRegistry {
    // ----- scheme identifiers (match AerePQCAttestation) -----
    uint8 public constant SCHEME_FALCON512 = 1;
    uint8 public constant SCHEME_FALCON1024 = 2;
    uint8 public constant SCHEME_MLDSA44 = 3;
    uint8 public constant SCHEME_SLHDSA128S = 4;

    // ----- live precompile addresses on chain 2800 -----
    address public constant PRECOMPILE_FALCON512 = address(0x0AE1);
    address public constant PRECOMPILE_FALCON1024 = address(0x0AE2);
    address public constant PRECOMPILE_MLDSA44 = address(0x0AE3);
    address public constant PRECOMPILE_SLHDSA128S = address(0x0AE4);

    // ----- public key sizes / headers -----
    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

    // ----- signature sizes (fixed-length schemes) -----
    uint256 internal constant MLDSA44_SIG_LEN = 2420;
    uint256 internal constant SLHDSA128S_SIG_LEN = 7856;
    uint256 internal constant FALCON_NONCE_LEN = 40;

    // ----- domain separator for the proof-of-possession challenge -----
    bytes32 public constant POP_DOMAIN = keccak256("AerePQCKeyRegistry.v1.pop");

    // ----- key lifecycle -----
    // NONE=0 (never registered), ACTIVE=1, ROTATED=2 (superseded by a successor),
    // REVOKED=3 (terminal). Only ACTIVE keys are "current"; ROTATED keys still verify
    // cryptographically (historical), REVOKED keys never verify.
    uint8 public constant STATUS_NONE = 0;
    uint8 public constant STATUS_ACTIVE = 1;
    uint8 public constant STATUS_ROTATED = 2;
    uint8 public constant STATUS_REVOKED = 3;

    uint256 internal constant NO_KEY = type(uint256).max; // sentinel: "no linked key"

    struct KeyRecord {
        address owner; // identity that proved possession and controls the key
        uint8 scheme; // one of SCHEME_*
        uint8 status; // one of STATUS_*
        uint64 registeredBlock; // block the key was registered at
        uint256 rotatedFromId; // predecessor keyId (NO_KEY if registered fresh)
        uint256 rotatedToId; // successor keyId (NO_KEY until rotated away)
        bytes pubKey; // the raw NIST public key for the scheme
    }

    KeyRecord[] private _keys; // keyId is the index into this array

    // Per-identity monotonic nonce making every proof-of-possession challenge unique
    // (prevents replay across registrations by the same identity).
    mapping(address => uint64) public identityNonce;

    // All keyIds ever registered by an identity (append-only; includes retired/revoked).
    mapping(address => uint256[]) private _identityKeys;

    // ----- events -----
    event KeyRegistered(uint256 indexed keyId, address indexed owner, uint8 indexed scheme, uint64 popNonce);
    event KeyRotated(uint256 indexed oldKeyId, uint256 indexed newKeyId, address indexed owner);
    event KeyRevoked(uint256 indexed keyId, address indexed owner);

    error InvalidScheme(uint8 scheme);
    error InvalidPubKey();
    error UnknownKey(uint256 keyId);
    error NotKeyOwner(uint256 keyId);
    error KeyNotActive(uint256 keyId);
    error KeyAlreadyRevoked(uint256 keyId);
    error ProofOfPossessionFailed();

    // ==========================================================================
    //                              Registration
    // ==========================================================================

    /**
     * @notice Register a PQC public key under msg.sender, proving possession on-chain.
     *         The proof is a valid PQC `signature` by `pubKey` over the current
     *         proof-of-possession challenge for (msg.sender, scheme, pubKey). On
     *         success the key is stored ACTIVE and the identity nonce advances so the
     *         same proof cannot be replayed.
     * @param scheme    1=Falcon-512, 2=Falcon-1024, 3=ML-DSA-44, 4=SLH-DSA-128s.
     * @param pubKey    the raw NIST public key (length/header validated per scheme).
     * @param signature the PQC proof-of-possession envelope (see contract NatSpec).
     * @return keyId    the identifier assigned to the new key.
     */
    function registerKey(uint8 scheme, bytes calldata pubKey, bytes calldata signature)
        external
        returns (uint256 keyId)
    {
        keyId = _registerWithPoP(msg.sender, scheme, pubKey, signature, NO_KEY);
    }

    /**
     * @notice Rotate `oldKeyId` to a fresh successor key, proving possession of the NEW
     *         key. Only the identity that owns `oldKeyId` may rotate it, and only while
     *         `oldKeyId` is ACTIVE. The old key becomes ROTATED (it still verifies as a
     *         historical key but is no longer current); the new key is ACTIVE and linked
     *         to its predecessor. The new key may use any of the four schemes.
     * @return newKeyId the identifier assigned to the successor key.
     */
    function rotateKey(uint256 oldKeyId, uint8 newScheme, bytes calldata newPubKey, bytes calldata signature)
        external
        returns (uint256 newKeyId)
    {
        if (oldKeyId >= _keys.length) revert UnknownKey(oldKeyId);
        KeyRecord storage old = _keys[oldKeyId];
        if (old.owner != msg.sender) revert NotKeyOwner(oldKeyId);
        if (old.status != STATUS_ACTIVE) revert KeyNotActive(oldKeyId);

        newKeyId = _registerWithPoP(msg.sender, newScheme, newPubKey, signature, oldKeyId);

        old.status = STATUS_ROTATED;
        old.rotatedToId = newKeyId;

        emit KeyRotated(oldKeyId, newKeyId, msg.sender);
    }

    /**
     * @notice Permanently revoke a key you own. Terminal: a revoked key can never
     *         verify again and cannot be reactivated. Callable while ACTIVE or ROTATED.
     */
    function revokeKey(uint256 keyId) external {
        if (keyId >= _keys.length) revert UnknownKey(keyId);
        KeyRecord storage k = _keys[keyId];
        if (k.owner != msg.sender) revert NotKeyOwner(keyId);
        if (k.status == STATUS_REVOKED) revert KeyAlreadyRevoked(keyId);

        k.status = STATUS_REVOKED;
        emit KeyRevoked(keyId, msg.sender);
    }

    // ==========================================================================
    //                         Internal registration core
    // ==========================================================================

    function _registerWithPoP(
        address owner,
        uint8 scheme,
        bytes calldata pubKey,
        bytes calldata signature,
        uint256 rotatedFromId
    ) internal returns (uint256 keyId) {
        _validatePubKey(scheme, pubKey);

        uint64 nonce = identityNonce[owner];
        bytes32 challenge = _popChallenge(owner, scheme, keccak256(pubKey), nonce);

        // Full on-chain PQC proof-of-possession over the contract-bound challenge.
        if (!_verify(scheme, pubKey, challenge, signature)) revert ProofOfPossessionFailed();

        // Effects: advance the identity nonce BEFORE storing so any reentrant view
        // sees a consumed proof, and the same signature can never be replayed.
        identityNonce[owner] = nonce + 1;

        keyId = _keys.length;
        _keys.push(
            KeyRecord({
                owner: owner,
                scheme: scheme,
                status: STATUS_ACTIVE,
                registeredBlock: uint64(block.number),
                rotatedFromId: rotatedFromId,
                rotatedToId: NO_KEY,
                pubKey: pubKey
            })
        );
        _identityKeys[owner].push(keyId);

        emit KeyRegistered(keyId, owner, scheme, nonce);
    }

    function _validatePubKey(uint8 scheme, bytes calldata pubKey) internal pure {
        if (scheme == SCHEME_FALCON512) {
            if (pubKey.length != FALCON512_PK_LEN || uint8(pubKey[0]) != FALCON512_PK_HEADER) revert InvalidPubKey();
        } else if (scheme == SCHEME_FALCON1024) {
            if (pubKey.length != FALCON1024_PK_LEN || uint8(pubKey[0]) != FALCON1024_PK_HEADER) revert InvalidPubKey();
        } else if (scheme == SCHEME_MLDSA44) {
            if (pubKey.length != MLDSA44_PK_LEN) revert InvalidPubKey();
        } else if (scheme == SCHEME_SLHDSA128S) {
            if (pubKey.length != SLHDSA128S_PK_LEN) revert InvalidPubKey();
        } else {
            revert InvalidScheme(scheme);
        }
    }

    // ==========================================================================
    //                          Challenges (views)
    // ==========================================================================

    /**
     * @notice The proof-of-possession challenge an identity must sign to register/rotate
     *         a key at a given nonce. Off-chain signers query this to know what to sign.
     */
    function popChallenge(address owner, uint8 scheme, bytes calldata pubKey, uint64 nonce)
        external
        view
        returns (bytes32)
    {
        return _popChallenge(owner, scheme, keccak256(pubKey), nonce);
    }

    /// @notice The challenge for `owner` to register `pubKey` right now (current nonce).
    function currentPopChallenge(address owner, uint8 scheme, bytes calldata pubKey) external view returns (bytes32) {
        return _popChallenge(owner, scheme, keccak256(pubKey), identityNonce[owner]);
    }

    function _popChallenge(address owner, uint8 scheme, bytes32 pubKeyHash, uint64 nonce)
        internal
        view
        returns (bytes32)
    {
        return keccak256(abi.encode(POP_DOMAIN, block.chainid, address(this), owner, scheme, pubKeyHash, nonce));
    }

    // ==========================================================================
    //                        Verification helpers (views)
    // ==========================================================================

    /**
     * @notice Verify a PQC signature over a raw 32-byte message with a REGISTERED key,
     *         via the live precompile. Fail-closed: returns false (never reverts) for an
     *         unknown key, a REVOKED key, or a malformed signature. ACTIVE and ROTATED
     *         keys verify (a ROTATED key is historically valid); REVOKED never does.
     *         Consumers that require a CURRENT key must also check statusOf()==ACTIVE.
     */
    function verifyWithKey(uint256 keyId, bytes32 message, bytes calldata signature) external view returns (bool) {
        if (keyId >= _keys.length) return false;
        KeyRecord storage k = _keys[keyId];
        if (k.status == STATUS_REVOKED || k.status == STATUS_NONE) return false;
        return _verify(k.scheme, k.pubKey, message, signature);
    }

    /**
     * @notice Stateless PQC verification for any scheme over a raw 32-byte message. No
     *         registration or binding: `message` is used directly as the signed message,
     *         so anyone can check any NIST PQC signature via eth_call for free.
     */
    function verify(uint8 scheme, bytes calldata pubKey, bytes32 message, bytes calldata signature)
        external
        view
        returns (bool)
    {
        return _verify(scheme, pubKey, message, signature);
    }

    // ==========================================================================
    //                            Internal verifier
    // ==========================================================================

    /// @dev Fail-closed: a malformed length or unknown scheme returns false (never
    ///      reverts), mirroring AereCryptoRegistry.verify. A well-formed-but-wrong
    ///      signature yields an input the precompile rejects with 0.
    function _verify(uint8 scheme, bytes memory pubKey, bytes32 message, bytes memory signature)
        internal
        view
        returns (bool)
    {
        (bool built, address precompile, bytes memory input) = _tryBuildInput(scheme, pubKey, message, signature);
        if (!built) return false;
        (bool ok, bytes memory ret) = precompile.staticcall(input);
        return ok && ret.length >= 32 && ret[31] == 0x01;
    }

    /// @dev Assemble (precompile, input) for a scheme (identical layout to
    ///      AerePQCAttestation). Returns built=false on any length/scheme that cannot form
    ///      a valid input, so verification fails closed rather than reverting.
    function _tryBuildInput(uint8 scheme, bytes memory pubKey, bytes32 message, bytes memory signature)
        internal
        pure
        returns (bool built, address precompile, bytes memory input)
    {
        if (scheme == SCHEME_FALCON512 || scheme == SCHEME_FALCON1024) {
            // signature = nonce(40) || esig
            if (signature.length <= FALCON_NONCE_LEN) return (false, address(0), "");
            uint256 sigLen = signature.length - FALCON_NONCE_LEN; // esig length
            if (sigLen > type(uint16).max) return (false, address(0), "");

            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;
        } else if (scheme == SCHEME_MLDSA44) {
            if (signature.length != MLDSA44_SIG_LEN) return (false, address(0), "");
            input = abi.encodePacked(pubKey, signature, message); // pk(1312) || sig(2420) || message(32)
            precompile = PRECOMPILE_MLDSA44;
        } else if (scheme == SCHEME_SLHDSA128S) {
            if (signature.length != SLHDSA128S_SIG_LEN) return (false, address(0), "");
            input = abi.encodePacked(pubKey, signature, message); // pk(32) || sig(7856) || message(32)
            precompile = PRECOMPILE_SLHDSA128S;
        } else {
            return (false, address(0), "");
        }
        built = true;
    }

    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];
        }
    }

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

    /// @notice Number of registered keys (keyIds run 0..keyCount-1).
    function keyCount() external view returns (uint256) {
        return _keys.length;
    }

    /// @notice Read a registered key. Reverts on an unknown keyId.
    function getKey(uint256 keyId)
        external
        view
        returns (
            address owner,
            uint8 scheme,
            uint8 status,
            uint64 registeredBlock,
            uint256 rotatedFromId,
            uint256 rotatedToId,
            bytes memory pubKey
        )
    {
        if (keyId >= _keys.length) revert UnknownKey(keyId);
        KeyRecord storage k = _keys[keyId];
        return (k.owner, k.scheme, k.status, k.registeredBlock, k.rotatedFromId, k.rotatedToId, k.pubKey);
    }

    /// @notice Lifecycle status of a key (STATUS_NONE for a never-registered id).
    function statusOf(uint256 keyId) external view returns (uint8) {
        if (keyId >= _keys.length) return STATUS_NONE;
        return _keys[keyId].status;
    }

    /// @notice True iff the key exists and is ACTIVE (current, not rotated/revoked).
    function isActiveKey(uint256 keyId) external view returns (bool) {
        if (keyId >= _keys.length) return false;
        return _keys[keyId].status == STATUS_ACTIVE;
    }

    /// @notice The scheme of a key, or 0 for an unknown id.
    function schemeOf(uint256 keyId) external view returns (uint8) {
        if (keyId >= _keys.length) return 0;
        return _keys[keyId].scheme;
    }

    /// @notice The address that owns/controls a key (zero for an unknown id).
    function ownerOf(uint256 keyId) external view returns (address) {
        if (keyId >= _keys.length) return address(0);
        return _keys[keyId].owner;
    }

    /// @notice All keyIds ever registered by `owner` (append-only, includes retired).
    function keysOf(address owner) external view returns (uint256[] memory) {
        return _identityKeys[owner];
    }

    /// @notice The live precompile address for a scheme (reverts on an unknown scheme).
    function precompileFor(uint8 scheme) external pure returns (address) {
        if (scheme == SCHEME_FALCON512) return PRECOMPILE_FALCON512;
        if (scheme == SCHEME_FALCON1024) return PRECOMPILE_FALCON1024;
        if (scheme == SCHEME_MLDSA44) return PRECOMPILE_MLDSA44;
        if (scheme == SCHEME_SLHDSA128S) return PRECOMPILE_SLHDSA128S;
        revert InvalidScheme(scheme);
    }
}
