// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

import {ECDSA} from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";

/// @title AereSessionKeyValidatorV2 - scoped session keys for AereModularAccount
/// @notice Audit-fix redeploy of AereSessionKeyValidator
///         (0x6e03A3D7A4c90d6f8dD6F0BA1F6e8aB1F8990D26, chain 2800). Same module
///         interface (ERC-7579 module type 1) so an account installs V2 purely by
///         address; the enableSession / onInstall / revokeSession / addPermission /
///         removePermission / sessions() / isValidSignatureWithSender surfaces are
///         byte-for-byte compatible with V1.
///
///         WHAT V1 GOT WRONG (CONFIRMED finding — per-batch value-cap bypass):
///         V1._checkCallData validated an executeBatch by running _checkSingle over
///         each Call INDEPENDENTLY, enforcing `value <= maxValuePerOp` per sub-call
///         with NO running sum. A session key could therefore move
///         maxValuePerOp * N native AERE in a single executeBatch (an N-way drain),
///         even though maxValuePerOp was meant to bound the whole userOp.
///
///         V2 FIXES + HARDENS:
///           (1) FIX  executeBatch now caps the SUM of calls[i].value at
///               maxValuePerOp. The per-op cap truly bounds the whole userOp; a
///               single execute is unchanged (its one value must be <= maxValuePerOp).
///           (2) DEFENSE-IN-DEPTH  a session-scoped sub-call whose target == the
///               account AND whose selector is installModule / uninstallModule /
///               setRootOwner / execute / executeBatch is ALWAYS rejected, regardless
///               of the allowlist. A session key can never re-enter the account to
///               rewrite ownership, (un)install modules, or nest execute/executeBatch
///               (re-entrant batch nesting is blocked by the execute/executeBatch
///               entries). This closes the class where an account mistakenly
///               allowlists (account, installModuleSelector): the account's own
///               execute() self-call would otherwise pass its onlyEntryPointSelfOrRoot
///               gate as msg.sender == address(this).
///           (3) BONUS  an OPTIONAL per-(account, sessionKey) rolling cumulative
///               spend cap `maxTotalValue` (0 = unlimited, the default) bounds TOTAL
///               native outflow authorized across ops, not just per-op. Backward
///               compatible: sessions created with the V1-shaped enableSession /
///               onInstall are uncapped exactly as before; the cap is opt-in via
///               enableSessionWithCap or setMaxTotalValue.
///
///         Canonical ERC-4337 validationData packing is preserved:
///           authorizer (160 bits) | validUntil (48) << 160 | validAfter (48) << 208
///         with authorizer = 0 on success, 1 on failure; AereModularAccount unpacks
///         it and enforces the time window because AereEntryPointV2 reverts on any
///         non-zero account return.
///
/// @dev    validateUserOp is state-MUTATING in V2 (V1's was view) so the rolling cap
///         can be accounted. It writes storage ONLY when a maxTotalValue cap is set;
///         uncapped sessions touch no storage, exactly like V1. The account invokes
///         it through the non-view IAereValidator interface, so this is a normal
///         (non-static) call and the write persists. Accounting is CONSERVATIVE: it
///         records the value a validated op is AUTHORIZED to move. Because AereEntry-
///         PointV2 runs the whole validation phase before execution and reverts the
///         entire handleOps on a validation failure, a rejected op writes nothing;
///         an op that passes validation but later reverts in execution can only
///         OVER-count, never under-count, so actual native outflow is always
///         bounded by maxTotalValue.
contract AereSessionKeyValidatorV2 {
    struct Session {
        bool enabled;
        uint48 validAfter;
        uint48 validUntil;
        uint256 maxValuePerOp;
    }

    uint256 internal constant SIG_VALIDATION_FAILED = 1;
    bytes4 internal constant EIP1271_MAGIC = 0x1626ba7e;
    bytes4 internal constant EIP1271_FAIL = 0xffffffff;

    /// @notice Sentinel selector matching an empty-calldata (native transfer) call.
    bytes4 public constant SELECTOR_NATIVE_TRANSFER = 0x00000000;

    /// @notice Selectors of AereModularAccount execution functions a session key
    ///         is allowed to enter through (as the OUTER userOp.callData).
    bytes4 public constant EXECUTE_SELECTOR = bytes4(keccak256("execute(address,uint256,bytes)"));
    bytes4 public constant EXECUTE_BATCH_SELECTOR = bytes4(keccak256("executeBatch((address,uint256,bytes)[])"));

    /// @notice Account-admin selectors a session key may NEVER call back INTO the
    ///         account (target == account) via a scoped sub-call. Public so the
    ///         guarantee is auditable on-chain.
    bytes4 public constant INSTALL_MODULE_SELECTOR = bytes4(keccak256("installModule(uint256,address,bytes)"));
    bytes4 public constant UNINSTALL_MODULE_SELECTOR = bytes4(keccak256("uninstallModule(uint256,address,bytes)"));
    bytes4 public constant SET_ROOT_OWNER_SELECTOR = bytes4(keccak256("setRootOwner(address)"));

    /// @dev Mirror of AereModularAccount.Call for callData decoding.
    struct Call {
        address target;
        uint256 value;
        bytes data;
    }

    /// @notice account => sessionKey => session parameters.
    mapping(address => mapping(address => Session)) public sessions;

    /// @notice account => sessionKey => target => selector => allowed.
    mapping(address => mapping(address => mapping(address => mapping(bytes4 => bool)))) public allowedCall;

    /// @notice account => sessionKey => rolling cumulative native-value cap
    ///         (0 = unlimited). Bounds TOTAL authorized outflow across ops.
    mapping(address => mapping(address => uint256)) public maxTotalValue;

    /// @notice account => sessionKey => cumulative authorized native value spent so
    ///         far under the active grant. Reset to 0 whenever the session is
    ///         (re-)enabled. Only advanced while a cap is set.
    mapping(address => mapping(address => uint256)) public spentTotal;

    event SessionEnabled(address indexed account, address indexed key, uint48 validAfter, uint48 validUntil, uint256 maxValuePerOp);
    event SessionRevoked(address indexed account, address indexed key);
    event PermissionAdded(address indexed account, address indexed key, address target, bytes4 selector);
    event PermissionRemoved(address indexed account, address indexed key, address target, bytes4 selector);
    event MaxTotalValueSet(address indexed account, address indexed key, uint256 maxTotalValue);

    error InvalidSessionKey();
    error InvalidWindow();
    error LengthMismatch();
    error SessionNotEnabled();

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

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

    /// @notice Optional initial session at install time. initData layout is
    ///         IDENTICAL to V1 (no maxTotalValue) so existing installers work
    ///         unchanged:
    ///         abi.encode(address key, uint48 validAfter, uint48 validUntil,
    ///                    uint256 maxValuePerOp, address[] targets, bytes4[] selectors)
    ///         Pass empty bytes to install with no session. A rolling cumulative cap
    ///         is added afterwards via setMaxTotalValue (or grant with
    ///         enableSessionWithCap instead of onInstall).
    function onInstall(bytes calldata data) external {
        if (data.length == 0) return;
        (
            address key,
            uint48 validAfter,
            uint48 validUntil,
            uint256 maxValuePerOp,
            address[] memory targets,
            bytes4[] memory selectors
        ) = abi.decode(data, (address, uint48, uint48, uint256, address[], bytes4[]));
        _enableSession(msg.sender, key, validAfter, validUntil, maxValuePerOp, 0, targets, selectors);
    }

    /// @notice Per-account session state is NOT bulk-cleared on uninstall (nested
    ///         mappings are not enumerable). The account removes routing on
    ///         uninstall, which makes stale state unreachable; revoke sessions
    ///         explicitly before a planned reinstall.
    function onUninstall(bytes calldata) external {}

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

    /// @notice V1-compatible grant (uncapped cumulative total). Drop-in identical.
    function enableSession(
        address key,
        uint48 validAfter,
        uint48 validUntil,
        uint256 maxValuePerOp,
        address[] calldata targets,
        bytes4[] calldata selectors
    ) external {
        address[] memory t = targets;
        bytes4[] memory s = selectors;
        _enableSession(msg.sender, key, validAfter, validUntil, maxValuePerOp, 0, t, s);
    }

    /// @notice V2 grant with an optional rolling cumulative cap (maxTotalValue,
    ///         0 = unlimited). Re-enabling resets spentTotal to 0.
    function enableSessionWithCap(
        address key,
        uint48 validAfter,
        uint48 validUntil,
        uint256 maxValuePerOp,
        uint256 maxTotalValue_,
        address[] calldata targets,
        bytes4[] calldata selectors
    ) external {
        address[] memory t = targets;
        bytes4[] memory s = selectors;
        _enableSession(msg.sender, key, validAfter, validUntil, maxValuePerOp, maxTotalValue_, t, s);
    }

    /// @notice Set / update the rolling cumulative cap for an existing session
    ///         (0 = unlimited). Does NOT reset spentTotal, so raising the ceiling
    ///         mid-grant keeps prior accounting; use enableSessionWithCap to start a
    ///         fresh budget.
    function setMaxTotalValue(address key, uint256 maxTotalValue_) external {
        if (!sessions[msg.sender][key].enabled) revert SessionNotEnabled();
        maxTotalValue[msg.sender][key] = maxTotalValue_;
        emit MaxTotalValueSet(msg.sender, key, maxTotalValue_);
    }

    function revokeSession(address key) external {
        if (!sessions[msg.sender][key].enabled) revert SessionNotEnabled();
        sessions[msg.sender][key].enabled = false;
        emit SessionRevoked(msg.sender, key);
    }

    function addPermission(address key, address target, bytes4 selector) external {
        if (!sessions[msg.sender][key].enabled) revert SessionNotEnabled();
        allowedCall[msg.sender][key][target][selector] = true;
        emit PermissionAdded(msg.sender, key, target, selector);
    }

    function removePermission(address key, address target, bytes4 selector) external {
        allowedCall[msg.sender][key][target][selector] = false;
        emit PermissionRemoved(msg.sender, key, target, selector);
    }

    function _enableSession(
        address account,
        address key,
        uint48 validAfter,
        uint48 validUntil,
        uint256 maxValuePerOp,
        uint256 maxTotalValue_,
        address[] memory targets,
        bytes4[] memory selectors
    ) internal {
        if (key == address(0)) revert InvalidSessionKey();
        if (validUntil != 0 && validUntil <= validAfter) revert InvalidWindow();
        if (targets.length != selectors.length) revert LengthMismatch();
        sessions[account][key] = Session({
            enabled: true,
            validAfter: validAfter,
            validUntil: validUntil,
            maxValuePerOp: maxValuePerOp
        });
        // A fresh grant starts a fresh cumulative budget.
        maxTotalValue[account][key] = maxTotalValue_;
        spentTotal[account][key] = 0;
        for (uint256 i = 0; i < targets.length; i++) {
            allowedCall[account][key][targets[i]][selectors[i]] = true;
            emit PermissionAdded(account, key, targets[i], selectors[i]);
        }
        emit SessionEnabled(account, key, validAfter, validUntil, maxValuePerOp);
        if (maxTotalValue_ != 0) emit MaxTotalValueSet(account, key, maxTotalValue_);
    }

    // ---- Validation (called by the account, msg.sender = account) -------------

    /// @param userOpHash       hash the session key signed (EIP-191 envelope).
    /// @param accountCallData  the userOp.callData; must be execute()/executeBatch()
    ///                         on the account, fully inside the session allowlist.
    /// @param signature        65-byte ECDSA signature by the session key.
    /// @dev State-mutating (see contract notes): advances spentTotal only when a
    ///      rolling cap is configured.
    function validateUserOp(
        bytes32 userOpHash,
        bytes calldata accountCallData,
        bytes calldata signature
    ) external returns (uint256 validationData) {
        address account = msg.sender;
        (address key, ECDSA.RecoverError err) =
            ECDSA.tryRecover(ECDSA.toEthSignedMessageHash(userOpHash), signature);
        if (err != ECDSA.RecoverError.NoError || key == address(0)) return SIG_VALIDATION_FAILED;

        Session memory s = sessions[account][key];
        if (!s.enabled) return SIG_VALIDATION_FAILED;

        (bool ok, uint256 opValue) = _checkCallData(account, key, s.maxValuePerOp, accountCallData);
        if (!ok) return SIG_VALIDATION_FAILED;

        // (3) rolling cumulative cap. Only touch storage when a cap is set, so
        // uncapped sessions behave exactly like V1 (no extra SSTORE).
        uint256 cap = maxTotalValue[account][key];
        if (cap != 0) {
            uint256 spent = spentTotal[account][key]; // invariant: spent <= cap
            // overflow-safe: check remaining headroom instead of adding first.
            if (opValue > cap - spent) return SIG_VALIDATION_FAILED;
            spentTotal[account][key] = spent + opValue;
        }

        // Canonical 4337 packing; the account enforces the time window.
        return (uint256(s.validUntil) << 160) | (uint256(s.validAfter) << 208);
    }

    /// @return ok        whether the whole callData is in scope + within caps.
    /// @return totalValue the native value the op is authorized to move (single
    ///                    value, or the batch sum).
    function _checkCallData(
        address account,
        address key,
        uint256 maxValuePerOp,
        bytes calldata accountCallData
    ) internal view returns (bool ok, uint256 totalValue) {
        if (accountCallData.length < 4) return (false, 0);
        bytes4 outer = bytes4(accountCallData[0:4]);

        if (outer == EXECUTE_SELECTOR) {
            (address target, uint256 value, bytes memory data) =
                abi.decode(accountCallData[4:], (address, uint256, bytes));
            if (value > maxValuePerOp) return (false, 0);
            if (!_scopeOk(account, key, target, data)) return (false, 0);
            return (true, value);
        }

        if (outer == EXECUTE_BATCH_SELECTOR) {
            Call[] memory calls = abi.decode(accountCallData[4:], (Call[]));
            if (calls.length == 0) return (false, 0);
            uint256 sum = 0;
            for (uint256 i = 0; i < calls.length; i++) {
                if (!_scopeOk(account, key, calls[i].target, calls[i].data)) return (false, 0);
                // (1) THE FIX: bound the SUM of sub-call values by maxValuePerOp.
                // Running check also subsumes the per-call cap (each value <= sum).
                sum += calls[i].value;
                if (sum > maxValuePerOp) return (false, 0);
            }
            return (true, sum);
        }

        // Session keys may only enter through execute/executeBatch.
        return (false, 0);
    }

    /// @dev Target + selector allowlist check, plus (2) the defense-in-depth guard
    ///      that a session key can never call an account-admin selector back into
    ///      its own account. Returns whether the (target, selector) pair is in scope.
    function _scopeOk(
        address account,
        address key,
        address target,
        bytes memory data
    ) internal view returns (bool) {
        bytes4 sel;
        if (data.length == 0) {
            sel = SELECTOR_NATIVE_TRANSFER;
        } else if (data.length >= 4) {
            assembly {
                sel := mload(add(data, 32))
            }
        } else {
            return false; // 1-3 byte calldata is never in scope
        }
        // (2) never re-enter the account to rewrite ownership / (un)install modules
        // / nest execute|executeBatch, regardless of what the allowlist says.
        if (target == account && _isForbiddenSelfSelector(sel)) return false;
        return allowedCall[account][key][target][sel];
    }

    function _isForbiddenSelfSelector(bytes4 sel) internal pure returns (bool) {
        return
            sel == INSTALL_MODULE_SELECTOR ||
            sel == UNINSTALL_MODULE_SELECTOR ||
            sel == SET_ROOT_OWNER_SELECTOR ||
            sel == EXECUTE_SELECTOR ||
            sel == EXECUTE_BATCH_SELECTOR;
    }

    // ---- EIP-1271 pass-through (msg.sender = account) --------------------------

    /// @notice Validates an off-chain signature by a currently-active session key.
    ///         Scoping does not apply (no call is being made); only key validity
    ///         and the time window are enforced.
    function isValidSignatureWithSender(address, bytes32 hash, bytes calldata signature)
        external
        view
        returns (bytes4)
    {
        address account = msg.sender;
        (address key, ECDSA.RecoverError err) =
            ECDSA.tryRecover(ECDSA.toEthSignedMessageHash(hash), signature);
        if (err != ECDSA.RecoverError.NoError || key == address(0)) return EIP1271_FAIL;
        Session memory s = sessions[account][key];
        if (!s.enabled) return EIP1271_FAIL;
        if (s.validAfter != 0 && block.timestamp < s.validAfter) return EIP1271_FAIL;
        if (s.validUntil != 0 && block.timestamp > s.validUntil) return EIP1271_FAIL;
        return EIP1271_MAGIC;
    }
}
