#version 460

#extension GL_GOOGLE_include_directive : require
#extension GL_EXT_buffer_reference : require
#extension GL_EXT_scalar_block_layout : require

#include "bindless.glsl"
#include "scene_data.glsl"

// Forward+ light culling.
// One workgroup per LIGHT_TILE_SIZE x LIGHT_TILE_SIZE screen tile:
//   1. find the min/max depth of the tile from the depth pre-pass
//   2. build the tile frustum in view space
//   3. test every light against it and write the visible ones to the tile's light list

layout (push_constant, scalar) uniform constants
{
    SceneDataBuffer sceneData;
    uvec2 screenSize;
    uint depthTexID;
    uint padding;
} pcs;

layout (local_size_x = LIGHT_TILE_SIZE, local_size_y = LIGHT_TILE_SIZE, local_size_z = 1) in;

#define THREADS_PER_TILE (LIGHT_TILE_SIZE * LIGHT_TILE_SIZE)

shared uint tileMinDepth;
shared uint tileMaxDepth;
shared uint tileLightCount;

// left, right, bottom, top side planes (view space, through the camera origin)
shared vec4 tilePlanes[4];
// view space distances along -Z
shared float tileNearDistance;
shared float tileFarDistance;

vec3 unprojectToView(vec2 ndc, float depth)
{
    vec4 p = pcs.sceneData.invProj * vec4(ndc, depth, 1.0);
    return p.xyz / p.w;
}

// Plane through the camera origin and two points, normal facing towards `inside`
vec4 makeSidePlane(vec3 a, vec3 b, vec3 inside)
{
    vec3 n = normalize(cross(a, b));
    if (dot(n, inside) < 0.0) {
        n = -n;
    }
    return vec4(n, 0.0);
}

bool sphereInsideTile(vec3 center, float radius)
{
    float distance = -center.z;
    if (distance + radius < tileNearDistance || distance - radius > tileFarDistance) {
        return false;
    }

    for (int i = 0; i < 4; i++) {
        if (dot(tilePlanes[i].xyz, center) + tilePlanes[i].w < -radius) {
            return false;
        }
    }
    return true;
}

void main()
{
    uint localIndex = gl_LocalInvocationIndex;
    uvec2 tileID = gl_WorkGroupID.xy;
    uint tileIndex = tileID.y * pcs.sceneData.tileCountX + tileID.x;
    uint tileOffset = tileIndex * LIGHT_TILE_STRIDE;

    if (localIndex == 0) {
        tileMinDepth = 0xFFFFFFFFu;
        tileMaxDepth = 0u;
        tileLightCount = 0u;
    }
    memoryBarrierShared();
    barrier();

    // 1. min / max depth
    uvec2 pixel = gl_GlobalInvocationID.xy;
    if (all(lessThan(pixel, pcs.screenSize))) {
        float depth = texelFetch(
            sampler2D(textures[pcs.depthTexID], samplers[NEAREST_SAMPLER_ID]),
            ivec2(pixel),
            0).r;

        // depth == 1.0 is the cleared far plane (sky), nothing there receives light.
        // Depth is always positive, so its bits can be compared as uints.
        if (depth < 1.0) {
            uint depthBits = floatBitsToUint(depth);
            atomicMin(tileMinDepth, depthBits);
            atomicMax(tileMaxDepth, depthBits);
        }
    }
    memoryBarrierShared();
    barrier();

    // 2. tile frustum
    bool tileHasGeometry = tileMinDepth != 0xFFFFFFFFu;

    if (localIndex == 0 && tileHasGeometry) {
        vec2 screenSize = vec2(pcs.screenSize);
        vec2 ndcMin = vec2(tileID * LIGHT_TILE_SIZE) / screenSize * 2.0 - 1.0;
        vec2 ndcMax = vec2((tileID + 1) * LIGHT_TILE_SIZE) / screenSize * 2.0 - 1.0;
        vec2 ndcCenter = (ndcMin + ndcMax) * 0.5;

        vec3 corner00 = unprojectToView(vec2(ndcMin.x, ndcMin.y), 1.0);
        vec3 corner10 = unprojectToView(vec2(ndcMax.x, ndcMin.y), 1.0);
        vec3 corner11 = unprojectToView(vec2(ndcMax.x, ndcMax.y), 1.0);
        vec3 corner01 = unprojectToView(vec2(ndcMin.x, ndcMax.y), 1.0);
        vec3 center = unprojectToView(ndcCenter, 1.0);

        tilePlanes[0] = makeSidePlane(corner01, corner00, center);
        tilePlanes[1] = makeSidePlane(corner10, corner11, center);
        tilePlanes[2] = makeSidePlane(corner00, corner10, center);
        tilePlanes[3] = makeSidePlane(corner11, corner01, center);

        tileNearDistance = -unprojectToView(ndcCenter, uintBitsToFloat(tileMinDepth)).z;
        tileFarDistance = -unprojectToView(ndcCenter, uintBitsToFloat(tileMaxDepth)).z;
    }
    memoryBarrierShared();
    barrier();

    // 3. cull lights, every thread takes a slice of the light list
    if (tileHasGeometry) {
        for (uint i = localIndex; i < pcs.sceneData.numLights; i += THREADS_PER_TILE) {
            Light light = pcs.sceneData.lights.data[i];

            bool visible = light.type == TYPE_DIRECTIONAL_LIGHT;
            if (!visible) {
                vec3 viewPos = (pcs.sceneData.view * vec4(light.position, 1.0)).xyz;
                visible = sphereInsideTile(viewPos, light.range);
            }

            if (visible) {
                uint slot = atomicAdd(tileLightCount, 1u);
                if (slot < MAX_LIGHTS_PER_TILE) {
                    pcs.sceneData.tileLightIndices.data[tileOffset + 1 + slot] = i;
                }
            }
        }
    }
    memoryBarrierShared();
    barrier();

    if (localIndex == 0) {
        pcs.sceneData.tileLightIndices.data[tileOffset] = min(tileLightCount, uint(MAX_LIGHTS_PER_TILE));
    }
}