#version 430 core

layout(local_size_x = 16, local_size_y = 16) in;

// state_curr: read-only; state_next: write-only
layout(r8ui,  binding = 0) readonly  uniform uimage2D state_curr;
layout(r8ui,  binding = 1) writeonly uniform uimage2D state_next;

// claim buffer to resolve conflicts
layout(r32ui, binding = 2) coherent uniform uimage2D claim;

uniform ivec2 u_gridSize;
uniform uint  u_frame;

const uint AIR  = 0u;
const uint SAND = 1u;

uint hash(uvec2 p, uint seed) {
    p = p * 1664525u + 1013904223u + seed;
    p.x ^= p.y >> 16;
    p.y ^= p.x >> 16;
    return p.x ^ p.y;
}

void main() {
    ivec2 pos = ivec2(gl_GlobalInvocationID.xy);
    if (pos.x >= u_gridSize.x || pos.y >= u_gridSize.y) {
        return;
    }

    uint self = imageLoad(state_curr, pos).r;
    if (self != SAND) {
        // AIR cells: do nothing (next is cleared to AIR on CPU)
        return;
    }

    ivec2 dest = pos;

    // Hard floor: bottom row does not move down
    if (pos.y < u_gridSize.y - 1) {
        // 1) Try straight down
        ivec2 down = pos + ivec2(0, 1);
        bool canDown = (imageLoad(state_curr, down).r == AIR);

        if (canDown) {
            dest = down;
        } else {
            // 2) Try diagonals
            bool canLeft  = false;
            bool canRight = false;

            ivec2 downLeft  = pos + ivec2(-1, 1);
            ivec2 downRight = pos + ivec2( 1, 1);

            if (downLeft.x >= 0 && downLeft.y < u_gridSize.y) {
                if (imageLoad(state_curr, downLeft).r == AIR) {
                    canLeft = true;
                }
            }
            if (downRight.x < u_gridSize.x && downRight.y < u_gridSize.y) {
                if (imageLoad(state_curr, downRight).r == AIR) {
                    canRight = true;
                }
            }

            // Flip which side is "first" based on frame parity
            bool flip = ((u_frame & 1u) != 0u);

            ivec2 firstDest, secondDest;
            bool  canFirst, canSecond;

            if (!flip) {
                // even frames: left is "first"
                firstDest  = downLeft;
                secondDest = downRight;
                canFirst   = canLeft;
                canSecond  = canRight;
            } else {
                // odd frames: right is "first"
                firstDest  = downRight;
                secondDest = downLeft;
                canFirst   = canRight;
                canSecond  = canLeft;
            }

            if (canFirst && canSecond) {
                // both free: pick randomly, but "first"/"second" swap each frame
                uint h = hash(uvec2(pos), u_frame);
                bool useFirst = (h & 1u) != 0u;
                dest = useFirst ? firstDest : secondDest;
            } else if (canFirst) {
                dest = firstDest;
            } else if (canSecond) {
                dest = secondDest;
            }
            // else: no diagonal move, dest stays = pos
        }
    }

    // Write into next state with claim resolution
    if (dest.x == pos.x && dest.y == pos.y) {
        // Not moving; just keep sand here
        imageStore(state_next, pos, uvec4(SAND, 0u, 0u, 0u));
    } else {
        // Moving to a new cell; try to claim it
        uint old = imageAtomicCompSwap(claim, dest, 0u, 1u);

        if (old == 0u) {
            // Won the slot: move
            imageStore(state_next, dest, uvec4(SAND, 0u, 0u, 0u));
            // Original pos stays AIR because state_next was cleared
        } else {
            // Lost the slot: stay put
            imageStore(state_next, pos, uvec4(SAND, 0u, 0u, 0u));
        }
    }
}