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

import "forge-std/Test.sol";
import { IERC8004IdentityRegistry } from "../../contracts/identity/IERC8004IdentityRegistry.sol";
import { ERC8004IdentityAdapterV3 } from "../../contracts/identity/ERC8004IdentityAdapterV3.sol";
import { CanonicalIdentityRegistryV3 } from "../../contracts/identity/CanonicalIdentityRegistryV3.sol";
import { MandateRegistryV3 } from "../../contracts/mandate/MandateRegistryV3.sol";
import { MandateExecutionAdapterV3 } from "../../contracts/mandate/MandateExecutionAdapterV3.sol";

contract MockIdentityRegistry is IERC8004IdentityRegistry {
    mapping(uint256 agentId => Agent agent) internal _agents;
    mapping(uint256 agentId => mapping(address operator => bool allowed)) internal _authorized;

    function setAgent(Agent calldata agent) external {
        _agents[agent.agentId] = agent;
    }

    function setAuthorizedOperator(uint256 agentId, address operator, bool allowed) external {
        _authorized[agentId][operator] = allowed;
    }

    function getAgent(uint256 agentId) external view returns (Agent memory) {
        return _agents[agentId];
    }

    function isAuthorizedOperator(uint256 agentId, address caller) external view returns (bool) {
        return _authorized[agentId][caller];
    }

    function isActive(uint256 agentId) external view returns (bool) {
        return _agents[agentId].active;
    }
}

contract MockERC1271Wallet {
    bytes4 internal constant MAGIC = 0x1626ba7e;
    bytes32 public approvedDigest;

    function approveDigest(bytes32 digest) external {
        approvedDigest = digest;
    }

    function registerIdentity(
        CanonicalIdentityRegistryV3 identities,
        uint256 agentId,
        address agentWallet,
        address operator,
        address beneficiary
    ) external {
        identities.register(agentId, agentWallet, operator, beneficiary);
    }

    function isValidSignature(bytes32 digest, bytes calldata) external view returns (bytes4) {
        return digest == approvedDigest ? MAGIC : bytes4(0xffffffff);
    }
}

contract MandateTargetV3Mock {
    uint256 public calls;
    uint256 public totalValue;

    function createBond() external payable {
        ++calls;
        totalValue += msg.value;
    }
}

/// @dev Deliberately minimal module. Production modules must authenticate their external caller
///      before forwarding that caller as `operator` to the registry.
contract MandateModuleV3Mock {
    function consume(
        MandateRegistryV3 registry,
        MandateRegistryV3.Mandate memory mandate,
        bytes32 action,
        address target,
        bytes4 selector,
        uint256 amount,
        address operator
    ) external returns (bytes32) {
        return registry.consume(mandate, action, target, selector, amount, operator);
    }
}

contract MandateRegistryV3Test is Test {
    uint256 internal constant PRINCIPAL_KEY = 0xA11CE;
    address internal principal = vm.addr(PRINCIPAL_KEY);
    address internal operator = address(0x0F3A7001);
    address internal beneficiary = address(0xBEEF1C1A);
    address internal owner = address(0xA0A0);
    address internal attacker = address(0xBAD);
    address internal target;
    bytes32 internal constant ACTION = keccak256("createBond");
    bytes4 internal constant SELECTOR = bytes4(keccak256("createBond()"));

    CanonicalIdentityRegistryV3 internal identities;
    ERC8004IdentityAdapterV3 internal identityAdapter;
    MandateRegistryV3 internal registry;
    MandateModuleV3Mock internal module;
    MandateExecutionAdapterV3 internal adapter;
    MandateTargetV3Mock internal mandateTarget;

    event ModuleAuthorizationUpdatedV3(address indexed module, bool allowed, address indexed updatedBy);

    function setUp() public {
        identities = new CanonicalIdentityRegistryV3();
        identityAdapter = new ERC8004IdentityAdapterV3(address(this), address(identities), address(identities).codehash);
        registry = new MandateRegistryV3(owner, address(identityAdapter), address(identityAdapter).codehash);
        module = new MandateModuleV3Mock();
        adapter = new MandateExecutionAdapterV3(address(registry));
        mandateTarget = new MandateTargetV3Mock();
        target = address(mandateTarget);
        _setIdentity(principal, operator, beneficiary);
    }

    function test_registersEOAMandateAndSnapshotsIdentityMetadata() public {
        MandateRegistryV3.Mandate memory mandate = _mandate(1);
        bytes32 mandateHash = _register(mandate, PRINCIPAL_KEY);

        MandateRegistryV3.MandateState memory stored = registry.state(mandateHash);
        assertTrue(stored.registered);
        assertFalse(stored.revoked);
        assertEq(stored.principal, principal);
        assertEq(stored.operator, operator);
        assertEq(stored.beneficiary, beneficiary);
        assertEq(stored.identityRegistry, address(identityAdapter));
        assertEq(stored.identityRegistryChainId, block.chainid);
        assertEq(stored.principalAgentId, 1);
        assertEq(stored.operatorAgentId, 2);
        assertEq(stored.nonce, 1);
        assertEq(stored.notBefore, mandate.notBefore);
        assertEq(stored.notAfter, mandate.notAfter);
        assertEq(registry.mandateForNonce(principal, 1), mandateHash);
        assertTrue(registry.isActive(mandateHash));
    }

    function test_registersERC1271PrincipalMandate() public {
        MockERC1271Wallet wallet = new MockERC1271Wallet();
        wallet.registerIdentity(identities, 3, address(wallet), operator, beneficiary);
        wallet.registerIdentity(identities, 4, operator, operator, beneficiary);

        MandateRegistryV3.Mandate memory mandate = _mandate(2);
        mandate.principal = address(wallet);
        mandate.principalAgentId = 3;
        mandate.operatorAgentId = 4;
        wallet.approveDigest(_digest(mandate));

        bytes32 mandateHash = registry.registerMandate(mandate, hex"1271");
        assertEq(registry.state(mandateHash).principal, address(wallet));
    }

    function test_arbitraryRevocationAttackIsRejected() public {
        bytes32 mandateHash = _register(_mandate(3), PRINCIPAL_KEY);

        vm.expectRevert(MandateRegistryV3.NotPrincipal.selector);
        vm.prank(attacker);
        registry.revokeMandate(mandateHash, principal);

        vm.expectRevert(abi.encodeWithSelector(MandateRegistryV3.PrincipalMismatch.selector, principal, attacker));
        vm.prank(attacker);
        registry.revokeMandate(mandateHash, attacker);

        assertFalse(registry.state(mandateHash).revoked);
    }

    function test_principalAndOwnerCanRevoke() public {
        bytes32 principalMandateHash = _register(_mandate(4), PRINCIPAL_KEY);
        vm.prank(principal);
        registry.revokeMandate(principalMandateHash);
        assertTrue(registry.state(principalMandateHash).revoked);

        bytes32 ownerMandateHash = _register(_mandate(5), PRINCIPAL_KEY);
        vm.prank(owner);
        registry.revokeMandate(ownerMandateHash);
        assertTrue(registry.state(ownerMandateHash).revoked);
    }

    function test_unauthorizedDirectConsumeCannotGriefMandate() public {
        MandateRegistryV3.Mandate memory mandate = _mandate(6);
        _register(mandate, PRINCIPAL_KEY);

        vm.expectRevert(MandateRegistryV3.LegacyConsumeDisabled.selector);
        vm.prank(attacker);
        registry.consume(mandate, ACTION, target, SELECTOR, 1, operator);
    }

    function test_allowlistedModuleConsumesForBoundOperatorOnly() public {
        MandateRegistryV3.Mandate memory mandate = _mandate(7);
        bytes32 mandateHash = _register(mandate, PRINCIPAL_KEY);
        _allowModule(true);

        bytes32 result = _execute(mandate, 40, operator);
        assertEq(result, mandateHash);
        assertEq(registry.state(mandateHash).spentSoFar, 40);

        vm.expectRevert(MandateRegistryV3.UnauthorizedCaller.selector);
        _execute(mandate, 1, attacker);
    }

    function test_expiredMandateCannotBeConsumedAndIsNotActive() public {
        MandateRegistryV3.Mandate memory mandate = _mandate(8);
        mandate.notAfter = uint64(block.timestamp + 1 hours);
        bytes32 mandateHash = _register(mandate, PRINCIPAL_KEY);
        _allowModule(true);

        vm.warp(mandate.notAfter);
        assertFalse(registry.isActive(mandateHash));
        vm.expectRevert(MandateRegistryV3.Expired.selector);
        _execute(mandate, 1, operator);
    }

    function test_enforcesSpendCapAndHourlyRateLimit() public {
        MandateRegistryV3.Mandate memory capMandate = _mandate(9);
        capMandate.spendCap = 100;
        capMandate.allowedActions[0] = _tuple(capMandate);
        _register(capMandate, PRINCIPAL_KEY);
        _allowModule(true);
        _execute(capMandate, 60, operator);
        vm.expectRevert(abi.encodeWithSelector(MandateRegistryV3.SpendCapExceeded.selector, 41, 40));
        _execute(capMandate, 41, operator);

        MandateRegistryV3.Mandate memory rateMandate = _mandate(10);
        rateMandate.rateLimitPerHour = 1;
        _register(rateMandate, PRINCIPAL_KEY);
        _execute(rateMandate, 0, operator);
        vm.expectRevert(MandateRegistryV3.RateLimitExceeded.selector);
        _execute(rateMandate, 0, operator);
    }

    function test_acceptsArraysAtExplicitMaximums() public {
        MandateRegistryV3.Mandate memory mandate = _mandate(11);
        mandate.allowedActions = _actions(registry.MAX_ALLOWED_ACTIONS());
        mandate.prohibitedActions = _actions(registry.MAX_PROHIBITED_ACTIONS());
        mandate.allowedContracts = _contracts(registry.MAX_ALLOWED_CONTRACTS());
        mandate.allowedSelectors = _selectors(registry.MAX_ALLOWED_SELECTORS());
        _register(mandate, PRINCIPAL_KEY);
    }

    function test_rejectsAllowedActionsAboveMaximum() public {
        MandateRegistryV3.Mandate memory mandate = _mandate(12);
        mandate.allowedActions = _actions(registry.MAX_ALLOWED_ACTIONS() + 1);
        vm.expectRevert(
            abi.encodeWithSelector(
                MandateRegistryV3.TooManyAllowedActions.selector,
                registry.MAX_ALLOWED_ACTIONS() + 1,
                registry.MAX_ALLOWED_ACTIONS()
            )
        );
        registry.registerMandate(mandate, "");
    }

    function test_rejectsProhibitedActionsAboveMaximum() public {
        MandateRegistryV3.Mandate memory mandate = _mandate(13);
        mandate.prohibitedActions = _actions(registry.MAX_PROHIBITED_ACTIONS() + 1);
        vm.expectRevert(
            abi.encodeWithSelector(
                MandateRegistryV3.TooManyProhibitedActions.selector,
                registry.MAX_PROHIBITED_ACTIONS() + 1,
                registry.MAX_PROHIBITED_ACTIONS()
            )
        );
        registry.registerMandate(mandate, "");
    }

    function test_rejectsAllowedContractsAboveMaximum() public {
        MandateRegistryV3.Mandate memory mandate = _mandate(14);
        mandate.allowedContracts = _contracts(registry.MAX_ALLOWED_CONTRACTS() + 1);
        vm.expectRevert(
            abi.encodeWithSelector(
                MandateRegistryV3.TooManyAllowedContracts.selector,
                registry.MAX_ALLOWED_CONTRACTS() + 1,
                registry.MAX_ALLOWED_CONTRACTS()
            )
        );
        registry.registerMandate(mandate, "");
    }

    function test_rejectsAllowedSelectorsAboveMaximum() public {
        MandateRegistryV3.Mandate memory mandate = _mandate(15);
        mandate.allowedSelectors = _selectors(registry.MAX_ALLOWED_SELECTORS() + 1);
        vm.expectRevert(
            abi.encodeWithSelector(
                MandateRegistryV3.TooManyAllowedSelectors.selector,
                registry.MAX_ALLOWED_SELECTORS() + 1,
                registry.MAX_ALLOWED_SELECTORS()
            )
        );
        registry.registerMandate(mandate, "");
    }

    function test_emptyScopesAndZeroCapsFailClosed() public {
        MandateRegistryV3.Mandate memory mandate = _mandate(23);
        mandate.allowedContracts = new address[](0);
        vm.expectRevert(
            abi.encodeWithSelector(
                MandateRegistryV3.TooManyAllowedContracts.selector, uint256(0), registry.MAX_ALLOWED_CONTRACTS()
            )
        );
        registry.registerMandate(mandate, "");

        mandate = _mandate(24);
        mandate.spendCap = 0;
        vm.expectRevert(MandateRegistryV3.InvalidScope.selector);
        registry.registerMandate(mandate, "");
    }

    function test_replayAndNonceReuseAreRejectedButNewNonceWorks() public {
        MandateRegistryV3.Mandate memory first = _mandate(16);
        bytes32 firstMandateHash = _register(first, PRINCIPAL_KEY);

        bytes memory firstSig = _sign(first, PRINCIPAL_KEY);
        vm.expectRevert(MandateRegistryV3.AlreadyRegistered.selector);
        registry.registerMandate(first, firstSig);

        MandateRegistryV3.Mandate memory sameNonceDifferentMandate = _mandate(16);
        sameNonceDifferentMandate.spendCap = 1;
        bytes memory sameNonceSig = _sign(sameNonceDifferentMandate, PRINCIPAL_KEY);
        vm.expectRevert(
            abi.encodeWithSelector(MandateRegistryV3.NonceAlreadyUsed.selector, principal, uint96(16), firstMandateHash)
        );
        registry.registerMandate(sameNonceDifferentMandate, sameNonceSig);

        MandateRegistryV3.Mandate memory newNonce = _mandate(17);
        _register(newNonce, PRINCIPAL_KEY);
    }

    function test_rejectsWrongProtocolTimeWindowAndIdentityBinding() public {
        MandateRegistryV3.Mandate memory badVersion = _mandate(18);
        badVersion.protocolVersion = "2.0.0";
        vm.expectRevert(MandateRegistryV3.UnsupportedProtocolVersion.selector);
        registry.registerMandate(badVersion, "");

        MandateRegistryV3.Mandate memory badTime = _mandate(19);
        badTime.notAfter = badTime.notBefore;
        vm.expectRevert(MandateRegistryV3.InvalidTimeWindow.selector);
        registry.registerMandate(badTime, "");

        MandateRegistryV3.Mandate memory badIdentity = _mandate(20);
        badIdentity.identityRegistryChainId = block.chainid + 1;
        vm.expectRevert(MandateRegistryV3.InvalidIdentityRegistry.selector);
        registry.registerMandate(badIdentity, "");

        MandateRegistryV3.Mandate memory mismatchedOperator = _mandate(21);
        mismatchedOperator.operator = attacker;
        vm.expectRevert(MandateRegistryV3.InvalidPrincipalIdentity.selector);
        registry.registerMandate(mismatchedOperator, "");
    }

    function test_onlyOwnerMayManageModuleAllowlist() public {
        vm.expectRevert(MandateRegistryV3.NotOwner.selector);
        registry.setModuleAuthorization(address(adapter), true);

        vm.prank(owner);
        vm.expectEmit(true, false, true, true);
        emit ModuleAuthorizationUpdatedV3(address(adapter), true, owner);
        registry.setModuleAuthorization(address(adapter), true);
        assertTrue(registry.allowedModules(address(adapter)));

        vm.prank(owner);
        vm.expectRevert(abi.encodeWithSelector(MandateRegistryV3.InvalidExecutionAdapter.selector, address(module)));
        registry.setModuleAuthorization(address(module), true);
    }

    function test_executionAdapterRejectsResourceBindingMismatchBeforeTargetCall() public {
        MandateRegistryV3.Mandate memory mandate = _mandate(22);
        bytes32 mandateHash = _register(mandate, PRINCIPAL_KEY);
        _allowModule(true);

        MandateExecutionAdapterV3.Execution memory execution = MandateExecutionAdapterV3.Execution({
            action: ACTION,
            resource: keccak256("different-resource"),
            target: target,
            amount: 1,
            data: abi.encodeWithSelector(SELECTOR)
        });
        vm.deal(operator, 1);
        vm.expectRevert(MandateRegistryV3.BadExecutionBinding.selector);
        vm.prank(operator);
        adapter.execute{ value: 1 }(abi.encode(mandate), execution);

        assertEq(registry.state(mandateHash).spentSoFar, 0);
        assertEq(mandateTarget.calls(), 0);
    }

    function test_signedTupleRejectsCrossCombinationAndTrailingCalldata() public {
        MandateRegistryV3.Mandate memory mandate = _mandate(25);
        _register(mandate, PRINCIPAL_KEY);
        _allowModule(true);

        MandateExecutionAdapterV3.Execution memory trailing = MandateExecutionAdapterV3.Execution({
            action: ACTION,
            resource: ACTION,
            target: target,
            amount: 0,
            data: abi.encodePacked(abi.encodeWithSelector(SELECTOR), bytes32(uint256(1)))
        });
        vm.expectRevert(abi.encodeWithSelector(MandateRegistryV3.ActionNotAllowed.selector, ACTION));
        vm.prank(operator);
        adapter.execute(abi.encode(mandate), trailing);

        MandateTargetV3Mock other = new MandateTargetV3Mock();
        mandate.allowedContracts = new address[](2);
        mandate.allowedContracts[0] = address(other) < target ? address(other) : target;
        mandate.allowedContracts[1] = address(other) < target ? target : address(other);
        mandate.nonce = 26;
        _register(mandate, PRINCIPAL_KEY);
        MandateExecutionAdapterV3.Execution memory crossed = MandateExecutionAdapterV3.Execution({
            action: ACTION,
            resource: ACTION,
            target: address(other),
            amount: 0,
            data: abi.encodeWithSelector(SELECTOR)
        });
        vm.expectRevert(abi.encodeWithSelector(MandateRegistryV3.ActionNotAllowed.selector, ACTION));
        vm.prank(operator);
        adapter.execute(abi.encode(mandate), crossed);
    }

    function test_unpinnedIdentityRegistryCannotAuthorizeMandate() public {
        MandateRegistryV3.Mandate memory mandate = _mandate(27);
        mandate.identityRegistry = address(identities);
        vm.expectRevert(MandateRegistryV3.InvalidIdentityRegistry.selector);
        registry.registerMandate(mandate, "");
    }

    function _allowModule(bool allowed) internal {
        vm.prank(owner);
        registry.setModuleAuthorization(address(adapter), allowed);
    }

    function _execute(MandateRegistryV3.Mandate memory mandate, uint256 amount, address caller)
        internal
        returns (bytes32)
    {
        MandateExecutionAdapterV3.Execution memory execution = MandateExecutionAdapterV3.Execution({
            action: ACTION,
            resource: ACTION,
            target: target,
            amount: amount,
            data: abi.encodeWithSelector(SELECTOR)
        });
        vm.deal(caller, amount);
        vm.prank(caller);
        (bytes32 mandateHash,) = adapter.execute{ value: amount }(abi.encode(mandate), execution);
        return mandateHash;
    }

    function _setIdentity(address principal_, address operator_, address beneficiary_) internal {
        vm.startPrank(principal_);
        identities.register(1, principal_, operator_, beneficiary_);
        identities.register(2, operator_, operator_, beneficiary_);
        vm.stopPrank();
    }

    function _mandate(uint96 nonce) internal view returns (MandateRegistryV3.Mandate memory mandate) {
        mandate.identityRegistryChainId = block.chainid;
        mandate.identityRegistry = address(identityAdapter);
        mandate.principalAgentId = 1;
        mandate.principal = principal;
        mandate.operatorAgentId = 2;
        mandate.operator = operator;
        mandate.beneficiary = beneficiary;
        mandate.allowedActions = _actions(1);
        mandate.prohibitedActions = new bytes32[](0);
        mandate.allowedContracts = new address[](1);
        mandate.allowedContracts[0] = target;
        mandate.allowedSelectors = new bytes4[](1);
        mandate.allowedSelectors[0] = SELECTOR;
        mandate.spendAsset = address(0);
        mandate.spendCap = 1_000_000 ether;
        mandate.rateLimitPerHour = 100;
        mandate.evidencePolicyId = bytes32(0);
        mandate.disputePolicyId = bytes32(0);
        mandate.notBefore = uint64(block.timestamp);
        mandate.notAfter = uint64(block.timestamp + 7 days);
        mandate.nonce = nonce;
        mandate.chainId = block.chainid;
        mandate.protocolVersion = "3.0.0";
        mandate.allowedActions[0] = _tuple(mandate);
    }

    function _tuple(MandateRegistryV3.Mandate memory mandate) internal view returns (bytes32) {
        return keccak256(
            abi.encode(ACTION, target, SELECTOR, mandate.spendCap, keccak256(abi.encodeWithSelector(SELECTOR)))
        );
    }

    function _register(MandateRegistryV3.Mandate memory mandate, uint256 signerKey) internal returns (bytes32) {
        return registry.registerMandate(mandate, _sign(mandate, signerKey));
    }

    function _digest(MandateRegistryV3.Mandate memory mandate) internal view returns (bytes32) {
        return keccak256(abi.encodePacked("\x19\x01", registry.DOMAIN_SEPARATOR(), registry.hashMandate(mandate)));
    }

    function _sign(MandateRegistryV3.Mandate memory mandate, uint256 signerKey) internal view returns (bytes memory) {
        (uint8 v, bytes32 r, bytes32 s) = vm.sign(signerKey, _digest(mandate));
        return abi.encodePacked(r, s, v);
    }

    function _actions(uint256 length) internal pure returns (bytes32[] memory values) {
        values = new bytes32[](length);
        for (uint256 i; i < length; ++i) {
            values[i] = bytes32(i + 1);
        }
    }

    function _contracts(uint256 length) internal pure returns (address[] memory values) {
        values = new address[](length);
        for (uint256 i; i < length; ++i) {
            // The loop is bounded to MAX_ALLOWED_CONTRACTS + 1 (33), so this cast is safe.
            // forge-lint: disable-next-line(unsafe-typecast)
            values[i] = address(uint160(i + 1));
        }
    }

    function _selectors(uint256 length) internal pure returns (bytes4[] memory values) {
        values = new bytes4[](length);
        for (uint256 i; i < length; ++i) {
            // The loop is bounded to MAX_ALLOWED_SELECTORS + 1 (33), so this cast is safe.
            // forge-lint: disable-next-line(unsafe-typecast)
            values[i] = bytes4(uint32(i + 1));
        }
    }
}