pragma solidity 0.8.20;

import {Test} from "forge-std/Test.sol";

import {Gatekeeper} from "../../src/Gatekeeper.sol";
import {Hashes} from "../../src/libraries/Hashes.sol";
import {Checkpoints} from "../../src/libraries/Checkpoints.sol";
import {IGatekeeper} from "../../src/interfaces/IGatekeeper.sol";
import {IStorageHistory} from "../../src/interfaces/IStorageHistory.sol";

contract MockStaking {
    GatekeeperVerification public gatekeeper;
    address public governor;

    constructor(uint256 existential) {
        gatekeeper = new GatekeeperVerification(existential, address(0));
        governor = msg.sender;
    }

    function ghost(bytes32 receiver, uint256 amount) external returns (uint256) {
        require(msg.sender == governor);
        return gatekeeper.ghost(receiver, amount);
    }

    function createNewGatekeeper(uint256 existential, address previousGatekeeper) external {
        require(msg.sender == governor);
        GatekeeperVerification newGatekeeper = new GatekeeperVerification(existential, address(gatekeeper));
        address storageHistory = IGatekeeper(previousGatekeeper).storageHistory();
        IStorageHistory(storageHistory).setOwner(address(newGatekeeper));
        gatekeeper = newGatekeeper;
    }
}

contract GatekeeperVerification is Gatekeeper {
    using Checkpoints for Checkpoints.Trace256;
    using Checkpoints for Checkpoints.Trace160;

    constructor(uint256 existential, address previousWeaver) Gatekeeper(existential, previousWeaver) {}

    function filledEntries(uint256 session) public view returns (uint256) {
        return _filledEntries[session];
    }

    function slotLengths(uint256 session, uint256 globalIndex) public view returns (uint160) {
        uint256 slotIndex = globalIndex % SLOTS;
        return _slotLenghts[session][slotIndex].latest();
    }

    function slotValues(uint256 session, uint256 globalIndex) public view returns (bytes32) {
        uint256 slotIndex = globalIndex % SLOTS;
        uint256 valueIndex = globalIndex / SLOTS;
        return _slotValues[session][slotIndex][valueIndex];
    }

    function treeNodesLatest(uint256 session, uint256 globalIndex) public view returns (bytes32) {
        uint256 slotIndex = globalIndex % SLOTS;
        return bytes32(_treeNodes[session][slotIndex].latest());
    }

    function computePreimage(uint256 globalIndex, uint256 a, bytes32 r) public pure returns (bytes32) {
        uint256 valueIndex = globalIndex / SLOTS;
        return _computeArgumentsHash(valueIndex, a, r);
    }

    function testVerify(
        bytes32[] calldata proof,
        bytes32[] calldata values,
        uint256 globalIndex,
        uint256 amount,
        bytes32 who
    ) public pure returns (bytes32) {
        if (proof.length != DEPTH) { return bytes32(0); }
        if (values.length == 0) { return bytes32(0); }

        uint256 slotIndex = globalIndex % SLOTS;
        uint256 valueIndex = globalIndex / SLOTS;
        if (valueIndex >= values.length) { return bytes32(0); }

        bytes32 computedHash = _computeArgumentsHash(valueIndex, amount, who);
        if (computedHash != values[valueIndex]) { return bytes32(0); }
        computedHash = bytes32(0);

        uint256 i;
        for (; i < values.length;) {
            bytes32 currHash = Hashes.efficientKeccak256(values[i]);
            computedHash = Hashes.efficientKeccak256(computedHash, currHash);
            unchecked { ++i; }
        }

        computedHash = Hashes.efficientKeccak256(computedHash);

        i = 0;
        for (; i < proof.length; ) {
            if (slotIndex % 2 == 0) {
                 computedHash = Hashes.efficientKeccak256(computedHash, proof[i]);
            } else {
                 computedHash = Hashes.efficientKeccak256(proof[i], computedHash);
            }
            slotIndex >>= 1;
            unchecked { ++i; }
        }

        return computedHash;
    }
}

contract GatekeeperWeaverTest is Test {
    address constant ALICE       = 0x0000000000000000000000000000000000000001;
    address constant BOB         = 0x0000000000000000000000000000000000000002;
    uint256 constant EXISTENTIAL = 1337;
    uint256 constant AMOUNT      = 1 * 1e7;

    MockStaking staking;
    GatekeeperVerification gatekeeper;

    function setUp() public {
        vm.prank(ALICE);
        staking = new MockStaking(EXISTENTIAL);
        gatekeeper = staking.gatekeeper();
    }

    function test_insertationWorksAsExpected() public {
        uint256 currentWeavingSession = gatekeeper.currentWeavingSession();
        uint256 maxCount = gatekeeper.ENTRIES();
        uint256 globalIndex;

        (bytes32[] memory whos, uint256[] memory amounts) = _prepareArrays(maxCount);

        for (uint256 i = 0; i < maxCount; i++) {
            if (i % 5 == 0) { vm.roll(block.number + 1); }
            globalIndex = _insertWithAssert(currentWeavingSession, amounts[i], whos[i]);
            assertTrue(_verifyProof(globalIndex, currentWeavingSession, block.number, amounts[i], whos[i]));
        }

        globalIndex = _insertWithAssert(currentWeavingSession, amounts[69], whos[69]);
        uint256 newSession = gatekeeper.currentWeavingSession();

        assertEq(currentWeavingSession + 1, newSession);
        vm.roll(block.number + 1337);

        assertTrue(_verifyProof(globalIndex, newSession, block.number, amounts[69], whos[69]));
        assertTrue(_verifyProof(0, currentWeavingSession, block.number, amounts[0], whos[0]));
        assertTrue(_verifyProof(69, currentWeavingSession, block.number, amounts[69], whos[69]));
        assertTrue(_verifyProof(420, currentWeavingSession, block.number, amounts[420], whos[420]));
        assertTrue(_verifyProof(1337, currentWeavingSession, block.number, amounts[1337], whos[1337]));
        assertTrue(_verifyProof(2047, currentWeavingSession, block.number, amounts[2047], whos[2047]));

        assertTrue(_verifyProof(0, currentWeavingSession, 100, amounts[0], whos[0]));
        assertTrue(_verifyProof(69, currentWeavingSession, 100, amounts[69], whos[69]));
        assertTrue(_verifyProof(420, currentWeavingSession, 100, amounts[420], whos[420]));

        assertFalse(_verifyProof(globalIndex, newSession, 100, amounts[69], whos[69]));
        assertFalse(_verifyProof(1337, currentWeavingSession, 100, amounts[1337], whos[1337]));
        assertFalse(_verifyProof(2047, currentWeavingSession, 100, amounts[2047], whos[2047]));

        vm.roll(block.number + 420);

        vm.prank(ALICE);
        staking.createNewGatekeeper(EXISTENTIAL, address(gatekeeper));
        gatekeeper = staking.gatekeeper();

        uint256 finalSession = gatekeeper.currentWeavingSession();

        for (uint256 i = 0; i < 69; i++) {
            if (i % 2 == 0) { vm.roll(block.number + 1); }
            vm.prank(ALICE);
            staking.ghost(whos[i], amounts[i]);
        }

        assertTrue(_verifyProof(globalIndex, newSession, block.number, amounts[69], whos[69]));
        assertTrue(_verifyProof(0, currentWeavingSession, block.number, amounts[0], whos[0]));
        assertTrue(_verifyProof(69, currentWeavingSession, block.number, amounts[69], whos[69]));
        assertTrue(_verifyProof(420, currentWeavingSession, block.number, amounts[420], whos[420]));
        assertTrue(_verifyProof(1337, currentWeavingSession, block.number, amounts[1337], whos[1337]));
        assertTrue(_verifyProof(2047, currentWeavingSession, block.number, amounts[2047], whos[2047]));

        assertTrue(_verifyProof(0, currentWeavingSession, 100, amounts[0], whos[0]));
        assertTrue(_verifyProof(69, currentWeavingSession, 100, amounts[69], whos[69]));
        assertTrue(_verifyProof(420, currentWeavingSession, 100, amounts[420], whos[420]));

        assertFalse(_verifyProof(globalIndex, newSession, 100, amounts[69], whos[69]));
        assertFalse(_verifyProof(1337, currentWeavingSession, 100, amounts[1337], whos[1337]));
        assertFalse(_verifyProof(2047, currentWeavingSession, 100, amounts[2047], whos[2047]));

        assertTrue(_verifyProof(0, finalSession, block.number, amounts[0], whos[0]));
        assertTrue(_verifyProof(34, finalSession, block.number, amounts[34], whos[34]));
        assertTrue(_verifyProof(35, finalSession, block.number, amounts[35], whos[35]));
    }

    function _verifyProof(
        uint256 globalIndex,
        uint256 session,
        uint256 atBlock,
        uint256 amount,
        bytes32 who
    ) private view returns (bool) {
        bytes32[] memory proof = gatekeeper.getProof(globalIndex, session, atBlock);
        bytes32[] memory values = gatekeeper.getSlotValues(globalIndex, session, atBlock);

        bytes32 root1 = gatekeeper.testVerify(proof, values, globalIndex, amount, who);
        (bytes32 root2, uint256 entries) = gatekeeper.getRoot(session, atBlock);
        assertTrue(entries <= gatekeeper.ENTRIES());

        return root1 == root2;
    }

    function _insertWithAssert(
        uint256 session,
        uint256 amount,
        bytes32 who
    ) private returns (uint256 globalIndex) {
        uint256 targetSlot = gatekeeper.filledEntries(session);
        bytes32 previousHash = gatekeeper.treeNodesLatest(session, targetSlot);
        uint256 prevEntries = gatekeeper.filledEntries(session);
        uint160 prevLength = gatekeeper.slotLengths(session, targetSlot);

        vm.prank(ALICE);
        globalIndex = staking.ghost(who, amount);

        if (session + 1 == gatekeeper.currentWeavingSession()) {
            assertEq(prevLength, gatekeeper.DEPTH());
            assertEq(gatekeeper.slotLengths(session + 1, 0), 1);
            assertEq(prevEntries, gatekeeper.ENTRIES());
            assertEq(gatekeeper.filledEntries(session + 1), 1);
            assertEq(session + 1, gatekeeper.currentWeavingSession());

            session += 1;
            previousHash = bytes32(0);
            prevLength = gatekeeper.slotLengths(session, 0);
        } else  {
            assertEq(prevLength + 1, gatekeeper.slotLengths(session, targetSlot));
            assertEq(prevEntries + 1, gatekeeper.filledEntries(session));
        }

        {
            bytes32 preimage1 = gatekeeper.slotValues(session, globalIndex);
            bytes32 preimage2 = gatekeeper.computePreimage(globalIndex, amount, who);
            assertEq(preimage1, preimage2);
        }

        {
            bytes32 nodeHash1 = gatekeeper.treeNodesLatest(session, globalIndex);
            bytes32 preimageHash = Hashes.efficientKeccak256(
                gatekeeper.computePreimage(globalIndex, amount, who)
            );
            bytes32 nodeHash2 = Hashes.efficientKeccak256(previousHash, preimageHash);
            assertEq(nodeHash1, nodeHash2);
        }
    }

    function _prepareArrays(uint256 count) private pure returns (bytes32[] memory, uint256[] memory) {
        bytes32[] memory whos = new bytes32[](count);
        uint256[] memory amounts = new uint256[](count);

        for (uint256 i = 0; i < count; i++) {
            whos[i] = keccak256(abi.encodePacked("user", i));
            amounts[i] = EXISTENTIAL + 1 + i;
        }

        return (whos, amounts);
    }
}