// SPDX-License-Identifier: MIT OR Apache-2.0
pragma solidity ^0.8.18;

import "utils/Test.sol";

contract Mock {
    uint256 state = 0;

    function numberA() public pure returns (uint256) {
        return 1;
    }

    function numberB() public pure returns (uint256) {
        return 2;
    }

    function add(uint256 a, uint256 b) public pure returns (uint256) {
        return a + b;
    }

    function pay(uint256 a) public payable returns (uint256) {
        return a;
    }

    function noReturnValue() public {
        // Does nothing of value, but also ensures that Solidity will 100%
        // generate an `extcodesize` check.
        state += 1;
    }
}

contract NestedMock {
    Mock private inner;

    constructor(Mock _inner) {
        inner = _inner;
    }

    function sum() public view returns (uint256) {
        return inner.numberA() + inner.numberB();
    }
}

contract NestedMockDelegateCall {
    Mock private inner;

    constructor(Mock _inner) {
        inner = _inner;
    }

    function sum() public returns (uint256) {
        (, bytes memory dataA) = address(inner).delegatecall(abi.encodeWithSelector(Mock.numberA.selector));
        (, bytes memory dataB) = address(inner).delegatecall(abi.encodeWithSelector(Mock.numberB.selector));
        return abi.decode(dataA, (uint256)) + abi.decode(dataB, (uint256));
    }
}

contract MockCallTest is Test {
    function testMockGetters() public {
        Mock target = new Mock();

        // pre-mock
        assertEq(target.numberA(), 1);
        assertEq(target.numberB(), 2);

        vm.mockCall(address(target), abi.encodeWithSelector(target.numberB.selector), abi.encode(10));

        // post-mock
        assertEq(target.numberA(), 1);
        assertEq(target.numberB(), 10);
    }

    function testMockNested() public {
        Mock inner = new Mock();
        NestedMock target = new NestedMock(inner);

        // pre-mock
        assertEq(target.sum(), 3);

        vm.mockCall(address(inner), abi.encodeWithSelector(inner.numberB.selector), abi.encode(9));

        // post-mock
        assertEq(target.sum(), 10);
    }

    // Ref: https://github.com/foundry-rs/foundry/issues/8066
    function testMockNestedDelegate() public {
        Mock inner = new Mock();
        NestedMockDelegateCall target = new NestedMockDelegateCall(inner);

        assertEq(target.sum(), 3);

        vm.mockCall(address(inner), abi.encodeWithSelector(inner.numberB.selector), abi.encode(9));

        assertEq(target.sum(), 10);
    }

    function testMockSelector() public {
        Mock target = new Mock();
        assertEq(target.add(5, 5), 10);

        vm.mockCall(address(target), abi.encodeWithSelector(target.add.selector), abi.encode(11));

        assertEq(target.add(5, 5), 11);
    }

    function testMockCalldata() public {
        Mock target = new Mock();
        assertEq(target.add(5, 5), 10);
        assertEq(target.add(6, 4), 10);

        vm.mockCall(address(target), abi.encodeWithSelector(target.add.selector, 5, 5), abi.encode(11));

        assertEq(target.add(5, 5), 11);
        assertEq(target.add(6, 4), 10);
    }

    function testClearMockedCalls() public {
        Mock target = new Mock();

        vm.mockCall(address(target), abi.encodeWithSelector(target.numberB.selector), abi.encode(10));

        assertEq(target.numberA(), 1);
        assertEq(target.numberB(), 10);

        vm.clearMockedCalls();

        assertEq(target.numberA(), 1);
        assertEq(target.numberB(), 2);
    }

    function testMockCallMultiplePartialMatch() public {
        Mock mock = new Mock();

        vm.mockCall(address(mock), abi.encodeWithSelector(mock.add.selector), abi.encode(10));
        vm.mockCall(address(mock), abi.encodeWithSelector(mock.add.selector, 2), abi.encode(20));
        vm.mockCall(address(mock), abi.encodeWithSelector(mock.add.selector, 2, 3), abi.encode(30));

        assertEq(mock.add(1, 2), 10);
        assertEq(mock.add(2, 2), 20);
        assertEq(mock.add(2, 3), 30);
    }

    function testMockCallWithValue() public {
        Mock mock = new Mock();

        vm.mockCall(address(mock), 10, abi.encodeWithSelector(mock.pay.selector), abi.encode(10));

        assertEq(mock.pay{value: 10}(1), 10);
        assertEq(mock.pay(1), 1);

        for (uint256 i = 0; i < 100; i++) {
            vm.mockCall(address(mock), i, abi.encodeWithSelector(mock.pay.selector), abi.encode(i * 2));
        }

        assertEq(mock.pay(1), 0);
        assertEq(mock.pay{value: 10}(1), 20);
        assertEq(mock.pay{value: 50}(1), 100);
    }

    function testMockCallWithValueTransfersBalance() public {
        Mock mock = new Mock();
        uint256 value = 10;
        vm.deal(address(this), value);

        vm.mockCall(address(mock), value, abi.encodeWithSelector(mock.pay.selector), abi.encode(10));

        assertEq(address(mock).balance, 0);
        assertEq(mock.pay{value: value}(1), 10);
        assertEq(address(mock).balance, value);
        assertEq(address(this).balance, 0);
    }

    function testMockCallWithValueTransfersPrankedSenderBalance() public {
        Mock mock = new Mock();
        address sender = address(0xBEEF);
        uint256 value = 10;
        vm.deal(address(this), 0);
        vm.deal(sender, value);

        vm.mockCall(address(mock), value, abi.encodeWithSelector(mock.pay.selector), abi.encode(10));

        vm.prank(sender);
        assertEq(mock.pay{value: value}(1), 10);
        assertEq(address(mock).balance, value);
        assertEq(address(this).balance, 0);
        assertEq(sender.balance, 0);
    }

    function testMockCallWithValueCalldataPrecedence() public {
        Mock mock = new Mock();

        vm.mockCall(address(mock), 10, abi.encodeWithSelector(mock.pay.selector), abi.encode(10));
        vm.mockCall(address(mock), abi.encodeWithSelector(mock.pay.selector, 2), abi.encode(2));

        assertEq(mock.pay{value: 10}(1), 10);
        assertEq(mock.pay{value: 10}(2), 2);
        assertEq(mock.pay(2), 2);
    }

    function testMockCallEmptyAccount() public {
        Mock mock = Mock(address(100));

        vm.mockCall(address(mock), abi.encodeWithSelector(mock.add.selector), abi.encode(10));
        vm.mockCall(address(mock), mock.noReturnValue.selector, abi.encode());

        assertEq(mock.add(1, 2), 10);
        mock.noReturnValue();
    }

    // Ref: https://github.com/foundry-rs/foundry/issues/10703
    function testMockCallEmptyAccountNoCodeInjection() public {
        Mock mock = Mock(address(100));

        vm.mockCall(address(mock), abi.encodeWithSelector(mock.add.selector), abi.encode(10), false);

        // Mocked calls that return data succeed while the account remains codeless.
        assertEq(mock.add(1, 2), 10);
        assertEq(address(mock).code.length, 0);

        // Unmocked calls to the empty account revert in the caller.
        vm.expectRevert();
        this.exposedCallNumberA(mock);
    }

    function testMockCallNoCodeInjectionNoReturnValue() public {
        Mock mock = Mock(address(100));

        vm.mockCall(address(mock), abi.encodeWithSelector(mock.noReturnValue.selector), abi.encode(), false);

        // Mocked calls to functions without return values still revert in the caller due to the
        // `extcodesize` check on the codeless account.
        vm.expectRevert();
        this.exposedCallNoReturnValue(mock);
    }

    function testMockCallEmptyAccountInjectCode() public {
        Mock mock = Mock(address(100));

        // `injectCode = true` behaves like the other overloads: code is injected into the empty
        // account so that mocked calls to functions without return values succeed as well.
        vm.mockCall(address(mock), abi.encodeWithSelector(mock.add.selector), abi.encode(10), true);
        vm.mockCall(address(mock), abi.encodeWithSelector(mock.noReturnValue.selector), abi.encode(), true);

        assertEq(mock.add(1, 2), 10);
        mock.noReturnValue();
        assertEq(address(mock).code.length, 1);
    }

    function exposedCallNumberA(Mock mock) external view returns (uint256) {
        return mock.numberA();
    }

    function exposedCallNoReturnValue(Mock mock) external {
        mock.noReturnValue();
    }
}

contract MockCallRevertTest is Test {
    error TestError(bytes msg);

    bytes constant ERROR_MESSAGE = "ERROR_MESSAGE";

    function testMockGettersRevert() public {
        Mock target = new Mock();

        // pre-mock
        assertEq(target.numberA(), 1);
        assertEq(target.numberB(), 2);

        vm.mockCallRevert(address(target), target.numberB.selector, ERROR_MESSAGE);

        // post-mock
        assertEq(target.numberA(), 1);
        try target.numberB() {
            revert();
        } catch (bytes memory err) {
            require(keccak256(err) == keccak256(ERROR_MESSAGE));
        }
    }

    function testMockRevertWithCustomError() public {
        Mock target = new Mock();

        assertEq(target.numberA(), 1);
        assertEq(target.numberB(), 2);

        bytes memory customError = abi.encodeWithSelector(TestError.selector, ERROR_MESSAGE);

        vm.mockCallRevert(address(target), abi.encodeWithSelector(target.numberB.selector), customError);

        assertEq(target.numberA(), 1);
        try target.numberB() {
            revert();
        } catch (bytes memory err) {
            require(keccak256(err) == keccak256(customError));
        }
    }

    function testMockNestedRevert() public {
        Mock inner = new Mock();
        NestedMock target = new NestedMock(inner);

        assertEq(target.sum(), 3);

        vm.mockCallRevert(address(inner), abi.encodeWithSelector(inner.numberB.selector), ERROR_MESSAGE);

        try target.sum() {
            revert();
        } catch (bytes memory err) {
            require(keccak256(err) == keccak256(ERROR_MESSAGE));
        }
    }

    function testMockCalldataRevert() public {
        Mock target = new Mock();
        assertEq(target.add(5, 5), 10);
        assertEq(target.add(6, 4), 10);

        vm.mockCallRevert(address(target), abi.encodeWithSelector(target.add.selector, 5, 5), ERROR_MESSAGE);

        assertEq(target.add(6, 4), 10);

        try target.add(5, 5) {
            revert();
        } catch (bytes memory err) {
            require(keccak256(err) == keccak256(ERROR_MESSAGE));
        }
    }

    function testClearMockRevertedCalls() public {
        Mock target = new Mock();

        vm.mockCallRevert(address(target), abi.encodeWithSelector(target.numberB.selector), ERROR_MESSAGE);

        vm.clearMockedCalls();

        assertEq(target.numberA(), 1);
        assertEq(target.numberB(), 2);
    }

    function testMockCallRevertPartialMatch() public {
        Mock mock = new Mock();

        vm.mockCallRevert(address(mock), abi.encodeWithSelector(mock.add.selector, 2), ERROR_MESSAGE);

        assertEq(mock.add(1, 2), 3);

        try mock.add(2, 3) {
            revert();
        } catch (bytes memory err) {
            require(keccak256(err) == keccak256(ERROR_MESSAGE));
        }
    }

    function testMockCallRevertWithValue() public {
        Mock mock = new Mock();
        uint256 value = 10;
        vm.deal(address(this), value);

        vm.mockCallRevert(address(mock), value, abi.encodeWithSelector(mock.pay.selector), ERROR_MESSAGE);

        assertEq(mock.pay(1), 1);
        assertEq(mock.pay(2), 2);

        uint256 initSenderBalance = address(this).balance;
        uint256 initTargetBalance = address(mock).balance;

        try mock.pay{value: value}(1) {
            revert();
        } catch (bytes memory err) {
            require(keccak256(err) == keccak256(ERROR_MESSAGE));
        }

        assertEq(address(this).balance, initSenderBalance);
        assertEq(address(mock).balance, initTargetBalance);
    }

    function testMockCallResetsMockCallRevert() public {
        Mock mock = new Mock();

        vm.mockCallRevert(address(mock), abi.encodeWithSelector(mock.add.selector), ERROR_MESSAGE);

        vm.mockCall(address(mock), abi.encodeWithSelector(mock.add.selector), abi.encode(5));
        assertEq(mock.add(2, 3), 5);
    }

    function testMockCallRevertResetsMockCall() public {
        Mock mock = new Mock();

        vm.mockCall(address(mock), abi.encodeWithSelector(mock.add.selector), abi.encode(5));
        assertEq(mock.add(2, 3), 5);

        vm.mockCallRevert(address(mock), abi.encodeWithSelector(mock.add.selector), ERROR_MESSAGE);

        try mock.add(2, 3) {
            revert();
        } catch (bytes memory err) {
            require(keccak256(err) == keccak256(ERROR_MESSAGE));
        }
    }

    function testMockCallRevertWithCall() public {
        Mock mock = new Mock();

        bytes memory customError = abi.encodeWithSelector(TestError.selector, ERROR_MESSAGE);

        vm.mockCallRevert(address(mock), abi.encodeWithSelector(mock.add.selector), customError);

        (bool success, bytes memory data) = address(mock).call(abi.encodeWithSelector(Mock.add.selector, 2, 3));
        assertEq(success, false);
        assertEq(data, customError);
    }

    function testMockCallEmptyAccount() public {
        Mock mock = Mock(address(100));

        vm.mockCallRevert(address(mock), abi.encodeWithSelector(mock.add.selector), ERROR_MESSAGE);

        try mock.add(2, 3) {
            revert();
        } catch (bytes memory err) {
            require(keccak256(err) == keccak256(ERROR_MESSAGE));
        }
    }
}
