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

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

/// @notice Minimal ERC-7579-style module interface used by AereModularAccount.
interface IAereModule {
    function onInstall(bytes calldata data) external;
    function onUninstall(bytes calldata data) external;
    function isModuleType(uint256 moduleTypeId) external view returns (bool);
}

/// @notice Validator module (module type 1) interface.
/// @dev Adapted from ERC-7579: the account strips its routing prefix and passes
///      the inner signature explicitly, together with the userOp callData so the
///      validator can scope what the op is allowed to do.
interface IAereValidator is IAereModule {
    function validateUserOp(
        bytes32 userOpHash,
        bytes calldata accountCallData,
        bytes calldata signature
    ) external returns (uint256 validationData);

    function isValidSignatureWithSender(
        address sender,
        bytes32 hash,
        bytes calldata signature
    ) external view returns (bytes4);
}

/// @notice Deposit-view subset of AereEntryPointV2.
interface IAereEntryPointDeposits {
    function balanceOf(address account) external view returns (uint256);
}

/// @title AereModularAccount - ERC-4337 smart account with ERC-7579-style modules
/// @notice Targets AERE's own AereEntryPointV2 (0x8D6f40598d552fF0Cb358b6012cF4227B86aF770,
///         chain 2800) EXACTLY:
///         - The EntryPoint keeps a strictly sequential per-sender nonce in its own
///           storage (no 2D nonce keys), so validator ROUTING CANNOT use a nonce key.
///           ROUTING CHOICE (documented per spec): the FIRST SIGNATURE BYTE routes:
///             0x00 -> root-owner ECDSA fallback. Layout: 0x00 || 65-byte ECDSA sig
///                     (EIP-191 personal_sign envelope over userOpHash, matching the
///                     AerePasskeyAccountV2 repo convention).
///             0x01 -> installed validator module. Layout:
///                     0x01 || 20-byte validator address || validator payload.
///         - AereEntryPointV2 REVERTS on ANY non-zero return from validateUserOp
///           (it does not unpack 4337 time bounds). Validators therefore return the
///           CANONICAL 4337 packing (authorizer | validUntil<<160 | validAfter<<208)
///           and THIS ACCOUNT unpacks it, enforces the time window against
///           block.timestamp itself, and returns strictly 0 (valid) or 1 (invalid).
///         - The EntryPoint executes ops by calling sender.call(op.callData), so
///           userOp.callData must encode execute()/executeBatch() on this account.
///         - Gas is charged from the account's EntryPoint deposit; validateUserOp
///           tops up only the shortfall from the account balance.
/// @dev Module types: 1 = validator (signature validation), 2 = executor (may call
///      executeFromExecutor and setRootOwner, e.g. the social recovery module).
contract AereModularAccount {
    struct PackedUserOperation {
        address sender;
        uint256 nonce;
        bytes initCode;
        bytes callData;
        bytes32 accountGasLimits;
        uint256 preVerificationGas;
        bytes32 gasFees;
        bytes paymasterAndData;
        bytes signature;
    }

    struct Call {
        address target;
        uint256 value;
        bytes data;
    }

    uint256 public constant MODULE_TYPE_VALIDATOR = 1;
    uint256 public constant MODULE_TYPE_EXECUTOR = 2;
    bytes4 internal constant EIP1271_MAGIC = 0x1626ba7e;
    bytes4 internal constant EIP1271_FAIL = 0xffffffff;

    /// @notice ERC-4337 EntryPoint allowed to call validateUserOp/execute. Set once.
    address public entryPoint;

    /// @notice Root owner EOA: fallback validation + direct admin. Replaceable by
    ///         self-call or an installed executor module (social recovery).
    address public rootOwner;

    /// @notice Factory that deployed this account (gates initialize).
    address public immutable factory;

    /// @dev moduleTypeId => module => installed.
    mapping(uint256 => mapping(address => bool)) internal _installed;

    event ModuleInstalled(uint256 indexed moduleTypeId, address indexed module);
    event ModuleUninstalled(uint256 indexed moduleTypeId, address indexed module);
    event RootOwnerChanged(address indexed previousOwner, address indexed newOwner);
    event Executed(address indexed target, uint256 value, bytes data);
    event ExecutedBatch(uint256 count);

    error Unauthorized();
    error AlreadyInitialized();
    error InvalidEntryPoint();
    error InvalidRootOwner();
    error UnsupportedModuleType(uint256 moduleTypeId);
    error InvalidModule(address module);
    error ModuleAlreadyInstalled(address module);
    error ModuleNotInstalled(address module);
    error CallFailed(bytes returnData);

    /// @dev Admin gate. NOTE: in addition to the EntryPoint and self-calls, the
    ///      current root owner may call admin functions DIRECTLY (documented
    ///      convenience so module management does not require a userOp round-trip;
    ///      the root owner already has full authority via the 0x00 signature path).
    modifier onlyEntryPointSelfOrRoot() {
        if (msg.sender != entryPoint && msg.sender != address(this) && msg.sender != rootOwner) {
            revert Unauthorized();
        }
        _;
    }

    constructor() {
        factory = msg.sender;
    }

    /// @notice One-shot initializer called by the factory at deploy time.
    function initialize(address _entryPoint, address _rootOwner) external {
        if (msg.sender != factory) revert Unauthorized();
        if (entryPoint != address(0)) revert AlreadyInitialized();
        if (_entryPoint == address(0)) revert InvalidEntryPoint();
        if (_rootOwner == address(0)) revert InvalidRootOwner();
        entryPoint = _entryPoint;
        rootOwner = _rootOwner;
        emit RootOwnerChanged(address(0), _rootOwner);
    }

    // ---- ERC-7579-style module management ------------------------------------

    function installModule(uint256 moduleTypeId, address module, bytes calldata initData)
        external
        onlyEntryPointSelfOrRoot
    {
        if (moduleTypeId != MODULE_TYPE_VALIDATOR && moduleTypeId != MODULE_TYPE_EXECUTOR) {
            revert UnsupportedModuleType(moduleTypeId);
        }
        if (module.code.length == 0) revert InvalidModule(module);
        if (_installed[moduleTypeId][module]) revert ModuleAlreadyInstalled(module);
        if (!IAereModule(module).isModuleType(moduleTypeId)) revert InvalidModule(module);
        _installed[moduleTypeId][module] = true;
        IAereModule(module).onInstall(initData);
        emit ModuleInstalled(moduleTypeId, module);
    }

    function uninstallModule(uint256 moduleTypeId, address module, bytes calldata deInitData)
        external
        onlyEntryPointSelfOrRoot
    {
        if (!_installed[moduleTypeId][module]) revert ModuleNotInstalled(module);
        _installed[moduleTypeId][module] = false;
        // Best effort deinit: a reverting module must never brick uninstall.
        (bool ok, ) = module.call(abi.encodeWithSelector(IAereModule.onUninstall.selector, deInitData));
        (ok);
        emit ModuleUninstalled(moduleTypeId, module);
    }

    function isModuleInstalled(uint256 moduleTypeId, address module, bytes calldata)
        external
        view
        returns (bool)
    {
        return _installed[moduleTypeId][module];
    }

    function supportsModule(uint256 moduleTypeId) external pure returns (bool) {
        return moduleTypeId == MODULE_TYPE_VALIDATOR || moduleTypeId == MODULE_TYPE_EXECUTOR;
    }

    function accountId() external pure returns (string memory) {
        return "aere.modular-account.1.0.0";
    }

    // ---- ERC-4337 validation ---------------------------------------------------

    /// @notice Called by AereEntryPointV2. Returns strictly 0 (valid) or 1 (invalid)
    ///         because this EntryPoint reverts on any non-zero validationData.
    function validateUserOp(
        PackedUserOperation calldata userOp,
        bytes32 userOpHash,
        uint256 missingAccountFunds
    ) external returns (uint256 validationData) {
        if (msg.sender != entryPoint) revert Unauthorized();
        validationData = _validateSignature(userOp, userOpHash);
        if (validationData == 0) {
            _payPrefund(missingAccountFunds);
        }
    }

    function _validateSignature(PackedUserOperation calldata userOp, bytes32 userOpHash)
        internal
        returns (uint256)
    {
        bytes calldata sig = userOp.signature;
        if (sig.length == 0) return 1;
        uint8 prefix = uint8(sig[0]);

        if (prefix == 0x00) {
            // Root-owner fallback: 0x00 || 65-byte ECDSA over EIP-191(userOpHash).
            if (sig.length != 66) return 1;
            (address rec, ECDSA.RecoverError err) =
                ECDSA.tryRecover(ECDSA.toEthSignedMessageHash(userOpHash), sig[1:66]);
            if (err != ECDSA.RecoverError.NoError || rec != rootOwner) return 1;
            return 0;
        }

        if (prefix == 0x01) {
            // Validator route: 0x01 || 20-byte validator address || payload.
            if (sig.length < 21) return 1;
            address validator = address(bytes20(sig[1:21]));
            if (!_installed[MODULE_TYPE_VALIDATOR][validator]) return 1;
            uint256 vd = IAereValidator(validator).validateUserOp(userOpHash, userOp.callData, sig[21:]);
            // Canonical 4337 packing from the validator:
            //   authorizer (160 bits) | validUntil (48 bits) << 160 | validAfter (48 bits) << 208
            // AereEntryPointV2 cannot unpack this, so enforce locally.
            if (address(uint160(vd)) != address(0)) return 1; // sig failure / no aggregator support
            uint48 validUntil = uint48(vd >> 160);
            uint48 validAfter = uint48(vd >> 208);
            if (validAfter != 0 && block.timestamp < validAfter) return 1;
            if (validUntil != 0 && block.timestamp > validUntil) return 1;
            return 0;
        }

        return 1;
    }

    /// @dev AereEntryPointV2 passes maxCost and later charges the account's EP
    ///      deposit for actual gas. Top up only the shortfall from our balance.
    function _payPrefund(uint256 missingAccountFunds) internal {
        if (missingAccountFunds == 0) return;
        uint256 dep = IAereEntryPointDeposits(entryPoint).balanceOf(address(this));
        if (dep >= missingAccountFunds) return;
        uint256 need = missingAccountFunds - dep;
        if (need > address(this).balance) need = address(this).balance;
        if (need == 0) return;
        // EntryPoint.receive() credits our own deposit (depositTo(msg.sender)).
        (bool ok, ) = payable(entryPoint).call{value: need}("");
        (ok); // EntryPoint reverts later if still underfunded
    }

    // ---- Execution ---------------------------------------------------------------

    function execute(address target, uint256 value, bytes calldata data)
        external
        onlyEntryPointSelfOrRoot
        returns (bytes memory result)
    {
        result = _call(target, value, data);
        emit Executed(target, value, data);
    }

    function executeBatch(Call[] calldata calls)
        external
        onlyEntryPointSelfOrRoot
        returns (bytes[] memory results)
    {
        results = new bytes[](calls.length);
        for (uint256 i = 0; i < calls.length; i++) {
            results[i] = _call(calls[i].target, calls[i].value, calls[i].data);
        }
        emit ExecutedBatch(calls.length);
    }

    /// @notice Execution entry for installed executor modules (module type 2).
    function executeFromExecutor(address target, uint256 value, bytes calldata data)
        external
        returns (bytes memory result)
    {
        if (!_installed[MODULE_TYPE_EXECUTOR][msg.sender]) revert Unauthorized();
        result = _call(target, value, data);
        emit Executed(target, value, data);
    }

    /// @notice Replace the root owner. Callable by a self-call (root rotating its
    ///         own key via execute) or by an installed executor module (social
    ///         recovery after guardian approval + timelock).
    function setRootOwner(address newOwner) external {
        if (msg.sender != address(this) && !_installed[MODULE_TYPE_EXECUTOR][msg.sender]) {
            revert Unauthorized();
        }
        if (newOwner == address(0)) revert InvalidRootOwner();
        emit RootOwnerChanged(rootOwner, newOwner);
        rootOwner = newOwner;
    }

    function _call(address target, uint256 value, bytes calldata data)
        internal
        returns (bytes memory)
    {
        (bool ok, bytes memory ret) = target.call{value: value}(data);
        if (!ok) revert CallFailed(ret);
        return ret;
    }

    // ---- EIP-1271 ------------------------------------------------------------------

    /// @notice EIP-1271 validation, routed like userOp signatures:
    ///         0x00 || 65-byte root ECDSA, or 0x01 || validator || payload.
    ///         The hash is domain-wrapped (account, chainid) to prevent cross-account
    ///         replay, matching the AerePasskeyAccountV2 repo convention.
    function isValidSignature(bytes32 hash, bytes calldata signature)
        external
        view
        returns (bytes4)
    {
        if (signature.length == 0) return EIP1271_FAIL;
        uint8 prefix = uint8(signature[0]);
        bytes32 wrapped = _wrap1271(hash);

        if (prefix == 0x00) {
            if (signature.length != 66) return EIP1271_FAIL;
            (address rec, ECDSA.RecoverError err) =
                ECDSA.tryRecover(ECDSA.toEthSignedMessageHash(wrapped), signature[1:66]);
            if (err != ECDSA.RecoverError.NoError || rec != rootOwner) return EIP1271_FAIL;
            return EIP1271_MAGIC;
        }

        if (prefix == 0x01 && signature.length >= 21) {
            address validator = address(bytes20(signature[1:21]));
            if (!_installed[MODULE_TYPE_VALIDATOR][validator]) return EIP1271_FAIL;
            return IAereValidator(validator).isValidSignatureWithSender(msg.sender, wrapped, signature[21:]);
        }

        return EIP1271_FAIL;
    }

    function _wrap1271(bytes32 hash) internal view returns (bytes32) {
        return keccak256(abi.encode(
            keccak256("AereAccount1271(address account,uint256 chainid,bytes32 hash)"),
            address(this),
            block.chainid,
            hash
        ));
    }

    receive() external payable {}
}