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";

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

    constructor(
        address staking,
        uint256 ghostedSupply,
        uint256 existential,
        uint48 deployedAt,
        uint104 amountIn,
        uint104 amountOut
    ) Gatekeeper(
        staking,
        ghostedSupply,
        existential,
        deployedAt,
        amountIn,
        amountOut
    ) {}

    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; }
        }

        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;

    GatekeeperVerification gatekeeper;

    function setUp() public {
        gatekeeper = new GatekeeperVerification(ALICE, 0, EXISTENTIAL, 0, 0, 0);
    }

    function test_insertationWorksAsExpected() public {
        uint256 currentSession = gatekeeper.currentSession();
        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(currentSession, amounts[i], whos[i]);
            assertTrue(_verifyProof(globalIndex, currentSession, block.number, amounts[i], whos[i]));
        }

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

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

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

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

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

    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 = gatekeeper.ghost(who, amount);

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

            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);
    }
}