#[compute]
#version 450

layout(local_size_x = 4, local_size_y = 4, local_size_z = 4) in;

layout(set = 0, binding = 0, std430) buffer DataBuffer { float sample_points[]; } voxels;
layout(set = 0, binding = 1) uniform Params { int world_size; float threshold; float u_time; } params;
layout(set = 0, binding = 2, std430) buffer SurfaceBuffer { float surface_points[]; } surface;
layout(set = 0, binding = 5, std430) buffer IndexBuffer { uint indices[]; } mesh_indices;
layout(set = 0, binding = 6, std430) buffer Counter { uint count; } index_count;

uint index3(uint x, uint y, uint z) {
    return x + y * params.world_size + z * params.world_size * params.world_size;
}

bool has_vertex(ivec3 p) {
    // Boundary check for the voxel itself
    if (any(lessThan(p, ivec3(0))) || any(greaterThanEqual(p, ivec3(params.world_size)))) return false;
    uint idx = index3(uint(p.x), uint(p.y), uint(p.z));
    float val = surface.surface_points[idx * 3u];
    return (val >= 0.0 && val <= float(params.world_size));
}

// Translated QUAD_POINTS from your GDScript
const ivec3 QUAD_OFFSETS[3][4] = ivec3[3][4](
    // X-Axis Edges
    ivec3[](ivec3(0,0,-1), ivec3(0,-1,-1), ivec3(0,-1,0), ivec3(0,0,0)),
    // Y-Axis Edges
    ivec3[](ivec3(0,0,-1), ivec3(0,0,0), ivec3(-1,0,0), ivec3(-1,0,-1)),
    // Z-Axis Edges
    ivec3[](ivec3(0,0,0), ivec3(0,-1,0), ivec3(-1,-1,0), ivec3(-1,0,0))
);

const ivec3 AXIS[3] = ivec3[](ivec3(1,0,0), ivec3(0,1,0), ivec3(0,0,1));


void main() {
    memoryBarrierBuffer();
    ivec3 id = ivec3(gl_GlobalInvocationID);
    if (any(greaterThanEqual(id, ivec3(params.world_size)))) return;

    uint idx = index3(id.x, id.y, id.z);


    for (int i = 0; i < 3; i++) {
        ivec3 neighbor_id = id + AXIS[i];

        if (any(greaterThanEqual(neighbor_id, ivec3(params.world_size)))) continue;

        float d1 = voxels.sample_points[idx];
        float d2 = voxels.sample_points[index3(neighbor_id.x, neighbor_id.y, neighbor_id.z)];

        if ((d1 < 0.0) != (d2 < 0.0) && d1 != -d2) {
            ivec3 p0 = id + QUAD_OFFSETS[i][0];
            ivec3 p1 = id + QUAD_OFFSETS[i][1];
            ivec3 p2 = id + QUAD_OFFSETS[i][2];
            ivec3 p3 = id + QUAD_OFFSETS[i][3];

            if (has_vertex(p0) && has_vertex(p1) && has_vertex(p2) && has_vertex(p3)) {
                uint v0 = index3(p0.x, p0.y, p0.z);
                uint v1 = index3(p1.x, p1.y, p1.z);
                uint v2 = index3(p2.x, p2.y, p2.z);
                uint v3 = index3(p3.x, p3.y, p3.z);

                uint start = atomicAdd(index_count.count, 6);

                if (d1 < 0.0) {
                    mesh_indices.indices[start + 0] = v0;
                    mesh_indices.indices[start + 1] = v1;
                    mesh_indices.indices[start + 2] = v2;
                    mesh_indices.indices[start + 3] = v0;
                    mesh_indices.indices[start + 4] = v2;
                    mesh_indices.indices[start + 5] = v3;
                } else {
                    // Reverse the order for the other side of the surface
                    mesh_indices.indices[start + 0] = v0;
                    mesh_indices.indices[start + 1] = v3;
                    mesh_indices.indices[start + 2] = v2;
                    mesh_indices.indices[start + 3] = v0;
                    mesh_indices.indices[start + 4] = v2;
                    mesh_indices.indices[start + 5] = v1;
                }
            }
        }
    }
}
