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

/// @notice Root-owner surface of AereModularAccount used by this module.
interface IAereModularAccountOwner {
    function rootOwner() external view returns (address);
    function setRootOwner(address newOwner) external;
}

/// @notice Subset of AerePQCKeyRegistry (chain 2800) this module depends on. Each guardian is a
///         NIST PQC public key registered under the registry with an on-chain PROOF-OF-POSSESSION,
///         so a guardian keyId is a key the guardian has already proven it controls. Verification
///         of every recovery leg is delegated to `verifyWithKey`, which runs the full PQC check via
///         AERE's live precompiles (Falcon-512 0x0AE1, Falcon-1024 0x0AE2, ML-DSA-44 0x0AE3,
///         SLH-DSA-128s 0x0AE4). `isActiveKey` is used so a rotated or revoked guardian key can
///         never authorize a recovery.
interface IAerePQCKeyRegistry {
    function verifyWithKey(uint256 keyId, bytes32 message, bytes calldata signature) external view returns (bool);
    function isActiveKey(uint256 keyId) external view returns (bool);
    function ownerOf(uint256 keyId) external view returns (address);
    function schemeOf(uint256 keyId) external view returns (uint8);
    function keyCount() external view returns (uint256);
}

/// @title  AerePQCSocialRecoveryModule - quantum-durable M-of-N recovery for AereModularAccount
/// @notice ERC-7579 executor module (module type 2). Recovers an account's root owner when the
///         primary key is lost, using an M-of-N committee of POST-QUANTUM guardians instead of
///         classical ECDSA guardian addresses. Each guardian is a NIST PQC public key
///         (Falcon-512 / Falcon-1024 / ML-DSA-44 / SLH-DSA-128s) registered with proof-of-
///         possession in AerePQCKeyRegistry; a recovery is authorized by >= `threshold` DISTINCT
///         guardians, each supplying a valid PQC signature over a domain-separated challenge that
///         binds (account, newOwner, round, chainid). Every leg is verified IN FULL, on-chain, by
///         the registry's live precompiles. The state change is held behind the same 48h timelock
///         as AereSocialRecoveryModule, giving the true owner a window to cancel a malicious
///         recovery.
///
/// @dev    Singleton: per-account state is keyed by the account address, exactly like
///         AereSocialRecoveryModule. Install as module type 2 on the account; executeRecovery then
///         calls account.setRootOwner(), which AereModularAccount permits for installed executors.
///
///         The t-of-n leg pattern is taken VERBATIM from AereThresholdAccount:
///           - a `Leg` is {guardianIndex, signature};
///           - distinctness is counted by guardian INDEX via a `seen` bitmask, so one key can never
///             fill two slots (duplicate guardian keyIds are additionally rejected at config time);
///           - malformed / duplicate / non-verifying legs are ignored, never counted, so an
///             attacker can never inflate the distinct-valid count;
///           - `round` (a per-account monotone nonce) is inside the challenge and advances on
///             execute/cancel, so a set of legs is single-use and can never be replayed.
///
///         HONEST SCOPE. Application/account layer only. This changes nothing about AERE consensus,
///         which remains classical ECDSA QBFT (7 Foundation validators, f=2). What is post-quantum
///         here is the GUARDIAN AUTHORIZATION: recovery approval survives a quantum adversary that
///         can forge ECDSA, because it is gated on NIST PQC signatures verified by the live
///         precompiles, not on secp256k1 guardian keys.
contract AerePQCSocialRecoveryModule {
    /// @notice Timelock between a threshold-approved recovery and the owner swap. Matches
    ///         AereSocialRecoveryModule.RECOVERY_DELAY so the two modules behave identically in time.
    uint256 public constant RECOVERY_DELAY = 48 hours;

    uint256 public constant MODULE_TYPE_EXECUTOR = 2;

    /// @notice The AerePQCKeyRegistry guardians are registered against (immutable; set at deploy).
    IAerePQCKeyRegistry public immutable registry;

    /// @notice Domain separator for the recovery challenge (distinct from any other AERE PQC domain,
    ///         so a signature made here can never be replayed on the key registry or a PQC account).
    bytes32 public constant RECOVERY_DOMAIN = keccak256("AerePQCSocialRecoveryModule.v1.recovery");

    uint256 internal constant MAX_GUARDIANS = 255; // guardianIndex is uint8; bitmask fits in uint256

    /// @notice One authorizing leg: which guardian (index into the account's guardian keyId array)
    ///         plus that guardian's PQC signature envelope over the recovery challenge.
    struct Leg {
        uint8 guardianIndex;
        bytes signature;
    }

    struct Recovery {
        address newOwner;    // proposed new root owner
        uint64 executeAfter; // scheduling time + RECOVERY_DELAY
        uint256 round;       // recoveryNonce snapshot the legs were signed for
        bool active;
    }

    /// @notice account => ordered guardian keyIds (position = guardianIndex used in a Leg).
    mapping(address => uint256[]) internal _guardianKeyIds;
    /// @notice account => keyId => is a guardian (membership + duplicate guard).
    mapping(address => mapping(uint256 => bool)) public isGuardianKey;
    /// @notice account => M in the M-of-N committee.
    mapping(address => uint256) public thresholdOf;
    /// @notice account => monotone recovery round id (replay barrier for the legs).
    mapping(address => uint256) public recoveryNonce;
    /// @notice account => currently pending recovery (if any).
    mapping(address => Recovery) public activeRecovery;

    event GuardianKeyAdded(address indexed account, uint256 indexed keyId, uint8 scheme);
    event GuardianKeyRemoved(address indexed account, uint256 indexed keyId);
    event ThresholdChanged(address indexed account, uint256 threshold);
    event RecoveryScheduled(address indexed account, uint256 indexed round, address indexed newOwner, uint64 executeAfter, uint256 signers);
    event RecoveryExecuted(address indexed account, uint256 indexed round, address indexed newOwner);
    event RecoveryCancelled(address indexed account, uint256 indexed round, address canceller);

    error BadRegistry();
    error AlreadyConfigured();
    error NotConfigured();
    error InvalidGuardianSet();
    error DuplicateGuardianKey(uint256 keyId);
    error GuardianKeyNotActive(uint256 keyId);
    error InvalidGuardianKey(uint256 keyId);
    error NotGuardianKey();
    error InvalidThreshold();
    error InvalidNewOwner();
    error RecoveryAlreadyActive();
    error NoActiveRecovery();
    error BelowThreshold(uint256 got, uint256 need);
    error TimelockNotElapsed(uint64 executeAfter, uint256 nowTimestamp);
    error NotOwner();

    constructor(address _registry) {
        if (_registry == address(0)) revert BadRegistry();
        registry = IAerePQCKeyRegistry(_registry);
    }

    // ---- ERC-7579-style module base -------------------------------------------

    function isModuleType(uint256 moduleTypeId) external pure returns (bool) {
        return moduleTypeId == MODULE_TYPE_EXECUTOR;
    }

    function name() external pure returns (string memory) {
        return "aere.pqc-social-recovery.1.0.0";
    }

    /// @notice initData layout: abi.encode(uint256[] guardianKeyIds, uint256 threshold).
    ///         Each keyId must be a distinct, currently ACTIVE registry key NOT owned by the account
    ///         itself (a self-owned guardian would collapse the non-custodial recovery guarantee).
    function onInstall(bytes calldata data) external {
        address account = msg.sender;
        if (_guardianKeyIds[account].length != 0) revert AlreadyConfigured();
        (uint256[] memory keyIds, uint256 threshold) = abi.decode(data, (uint256[], uint256));
        uint256 n = keyIds.length;
        if (n == 0 || n > MAX_GUARDIANS) revert InvalidGuardianSet();
        for (uint256 i = 0; i < n; i++) {
            _addGuardianKey(account, keyIds[i]);
        }
        if (threshold == 0 || threshold > n) revert InvalidThreshold();
        thresholdOf[account] = threshold;
        emit ThresholdChanged(account, threshold);
    }

    /// @notice Full per-account cleanup so a later reinstall starts clean. `recoveryNonce` is
    ///         intentionally NOT reset: keeping it monotone across (un)installs means legs signed
    ///         for a prior round can never be replayed after a reinstall.
    function onUninstall(bytes calldata) external {
        address account = msg.sender;
        uint256[] storage list = _guardianKeyIds[account];
        for (uint256 i = 0; i < list.length; i++) {
            isGuardianKey[account][list[i]] = false;
        }
        delete _guardianKeyIds[account];
        delete thresholdOf[account];
        delete activeRecovery[account];
    }

    // ---- Owner-only guardian management (msg.sender == the smart account) ------

    function addGuardianKey(uint256 keyId) external {
        _addGuardianKey(msg.sender, keyId);
    }

    function removeGuardianKey(uint256 keyId) external {
        address account = msg.sender;
        if (!isGuardianKey[account][keyId]) revert NotGuardianKey();
        isGuardianKey[account][keyId] = false;
        uint256[] storage list = _guardianKeyIds[account];
        for (uint256 i = 0; i < list.length; i++) {
            if (list[i] == keyId) {
                list[i] = list[list.length - 1];
                list.pop();
                break;
            }
        }
        // Keep threshold satisfiable and the set non-empty.
        if (list.length == 0 || thresholdOf[account] > list.length) revert InvalidThreshold();
        emit GuardianKeyRemoved(account, keyId);
    }

    function setThreshold(uint256 threshold) external {
        address account = msg.sender;
        if (threshold == 0 || threshold > _guardianKeyIds[account].length) revert InvalidThreshold();
        thresholdOf[account] = threshold;
        emit ThresholdChanged(account, threshold);
    }

    function _addGuardianKey(address account, uint256 keyId) internal {
        if (isGuardianKey[account][keyId]) revert DuplicateGuardianKey(keyId);
        // The key must EXIST and be currently ACTIVE (a rotated/revoked key can never be a guardian).
        if (!registry.isActiveKey(keyId)) revert GuardianKeyNotActive(keyId);
        // A guardian key owned by the account itself would let the account's own key approve its own
        // recovery, defeating the "distinct third parties" guarantee.
        if (registry.ownerOf(keyId) == account) revert InvalidGuardianKey(keyId);
        if (_guardianKeyIds[account].length >= MAX_GUARDIANS) revert InvalidGuardianSet();
        isGuardianKey[account][keyId] = true;
        _guardianKeyIds[account].push(keyId);
        emit GuardianKeyAdded(account, keyId, registry.schemeOf(keyId));
    }

    // ---- Recovery flow ---------------------------------------------------------

    /// @notice The exact 32-byte challenge each guardian must PQC-sign to authorize replacing
    ///         `account`'s root owner with `newOwner` at `round`. Off-chain guardians query this
    ///         (or reconstruct it) to know what to sign. Binds account + newOwner + round + chainid
    ///         + this module, so a leg is scoped to one account, one target owner, one round, one
    ///         chain, and one module deployment.
    function recoveryChallenge(address account, address newOwner, uint256 round) public view returns (bytes32) {
        return keccak256(abi.encode(RECOVERY_DOMAIN, block.chainid, address(this), account, newOwner, round));
    }

    /// @notice The challenge for the round a scheduleRecovery call would use RIGHT NOW.
    function currentRecoveryChallenge(address account, address newOwner) external view returns (bytes32) {
        return recoveryChallenge(account, newOwner, recoveryNonce[account]);
    }

    /// @notice Schedule a recovery to `newOwner`. Permissionless: the PQC legs ARE the authorization
    ///         (a relayer can submit them), exactly like AereThresholdAccount.executeThreshold. The
    ///         legs must carry >= threshold DISTINCT valid guardian signatures over the current
    ///         round's challenge; below threshold reverts and nothing is scheduled. On success the
    ///         owner swap is armed for `block.timestamp + RECOVERY_DELAY`.
    function scheduleRecovery(address account, address newOwner, Leg[] calldata legs)
        external
        returns (uint256 round)
    {
        if (newOwner == address(0)) revert InvalidNewOwner();
        uint256 t = thresholdOf[account];
        if (t == 0) revert NotConfigured();
        if (activeRecovery[account].active) revert RecoveryAlreadyActive();

        round = recoveryNonce[account];
        bytes32 challenge = recoveryChallenge(account, newOwner, round);
        uint256 signers = _countLegs(account, challenge, legs);
        if (signers < t) revert BelowThreshold(signers, t);

        uint64 executeAfter = uint64(block.timestamp + RECOVERY_DELAY);
        activeRecovery[account] = Recovery({newOwner: newOwner, executeAfter: executeAfter, round: round, active: true});
        emit RecoveryScheduled(account, round, newOwner, executeAfter, signers);
    }

    /// @notice Execute a scheduled recovery once the 48h timelock elapses. Permissionless (anyone
    ///         may poke it, the authorization was proven at schedule time). Advances `recoveryNonce`
    ///         so the scheduling round's legs are single-use, then calls account.setRootOwner().
    function executeRecovery(address account) external {
        Recovery memory r = activeRecovery[account];
        if (!r.active) revert NoActiveRecovery();
        if (block.timestamp < r.executeAfter) revert TimelockNotElapsed(r.executeAfter, block.timestamp);

        // effects before interaction
        recoveryNonce[account] = r.round + 1;
        delete activeRecovery[account];

        IAereModularAccountOwner(account).setRootOwner(r.newOwner);
        emit RecoveryExecuted(account, r.round, r.newOwner);
    }

    /// @notice Cancel a pending recovery. Callable by the current root owner or the account itself
    ///         (the true owner's window to stop a malicious recovery). Advances `recoveryNonce` so
    ///         the cancelled round's legs can never be reused.
    function cancelRecovery(address account) external {
        if (msg.sender != account && msg.sender != IAereModularAccountOwner(account).rootOwner()) {
            revert NotOwner();
        }
        Recovery memory r = activeRecovery[account];
        if (!r.active) revert NoActiveRecovery();
        recoveryNonce[account] = r.round + 1;
        delete activeRecovery[account];
        emit RecoveryCancelled(account, r.round, msg.sender);
    }

    // ---- Threshold counting (t-of-n pattern from AereThresholdAccount) ----------

    /// @dev Count DISTINCT valid guardian signers over `challenge`. A leg is counted only if it is
    ///      in-range, not-yet-seen, references a currently ACTIVE guardian key, and verifies in full
    ///      via the registry precompile. Everything else (out-of-range index, duplicate index,
    ///      rotated/revoked key, malformed or wrong signature) is ignored, so the distinct-valid
    ///      count can never be inflated.
    function _countLegs(address account, bytes32 challenge, Leg[] calldata legs)
        internal
        view
        returns (uint256 valid)
    {
        uint256[] storage keyIds = _guardianKeyIds[account];
        uint256 n = keyIds.length;
        uint256 seen = 0;
        for (uint256 i = 0; i < legs.length; i++) {
            uint256 idx = legs[i].guardianIndex;
            if (idx >= n) continue;
            uint256 bit = uint256(1) << idx;
            if (seen & bit != 0) continue; // never double-count a guardian
            uint256 keyId = keyIds[idx];
            // A rotated/revoked guardian key must never authorize (verifyWithKey alone would still
            // accept a ROTATED key as historically valid, so gate on isActiveKey too).
            if (!registry.isActiveKey(keyId)) continue;
            if (registry.verifyWithKey(keyId, challenge, legs[i].signature)) {
                seen |= bit;
                valid += 1;
            }
        }
    }

    // ---- Views -----------------------------------------------------------------

    /// @notice The ordered guardian keyIds for `account` (position = guardianIndex used in a Leg).
    function guardianKeysOf(address account) external view returns (uint256[] memory) {
        return _guardianKeyIds[account];
    }

    function guardianCount(address account) external view returns (uint256) {
        return _guardianKeyIds[account].length;
    }

    /// @notice keyId at a given guardian index (reverts on out-of-range).
    function guardianKeyAt(address account, uint256 index) external view returns (uint256) {
        return _guardianKeyIds[account][index];
    }

    /// @notice Compact committee view: (threshold, size, round, pending?).
    function committee(address account)
        external
        view
        returns (uint256 threshold, uint256 size, uint256 round, bool pending)
    {
        return (thresholdOf[account], _guardianKeyIds[account].length, recoveryNonce[account], activeRecovery[account].active);
    }
}
