#version 450

// GPU Light Culling Compute Shader
// Implements tile-based deferred lighting (Forward+)
// Each tile gets a list of affecting lights

// ISSUE 20 FIX: Specialization constants for adaptive workgroup sizing
// Default 16x16, can be overridden at pipeline creation based on light count
layout(constant_id = 0) const uint WORKGROUP_SIZE_X = 16;
layout(constant_id = 1) const uint WORKGROUP_SIZE_Y = 16;

layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z = 1) in;

// Maximum lights per tile
#define MAX_LIGHTS_PER_TILE 256

// Light structure
struct Light {
    vec4 position;   // xyz = position, w = radius
    vec4 color;      // rgb = color, a = intensity
    vec4 direction;  // xyz = direction (for spot), w = type (0=point, 1=spot, 2=directional)
    vec4 params;     // x = innerConeAngle, y = outerConeAngle, z = falloff, w = enabled
};

// Push constants for screen info
layout(push_constant) uniform PushConstants {
    uvec2 screenSize;
    uint lightCount;
    uint _padding;
} pc;

// Input: All scene lights
layout(set = 0, binding = 0, std430) readonly buffer LightBuffer {
    Light lights[];
};

// Input: Depth buffer for frustum culling
layout(set = 0, binding = 1) uniform sampler2D depthBuffer;

// Output: Light indices per tile
// Format: [tileCount][MAX_LIGHTS_PER_TILE + 1]
// First element of each tile = light count for that tile
layout(set = 0, binding = 2, std430) writeonly buffer TileLightIndices {
    uint tileData[];
};

// View-projection matrix for frustum construction
layout(set = 0, binding = 3) uniform CameraData {
    mat4 view;
    mat4 projection;
    mat4 invProjection;
    vec4 cameraPos;
} camera;

// Shared memory for light indices within workgroup
shared uint sharedLightCount;
shared uint sharedLightIndices[MAX_LIGHTS_PER_TILE];
shared uint sharedMinDepthInt;
shared uint sharedMaxDepthInt;
// OPTIMIZATION: Store frustum in shared memory so only thread 0 computes it
shared vec4 sharedFrustumPlanes[6];

// Convert screen UV to view-space position at given depth
vec3 screenToView(vec2 uv, float depth) {
    vec4 clip = vec4(uv * 2.0 - 1.0, depth, 1.0);
    vec4 view = camera.invProjection * clip;
    return view.xyz / view.w;
}

// Construct frustum planes for a tile
void getTileFrustum(uvec2 tileId, float minDepth, float maxDepth, out vec4 planes[6]) {
    // Tile corners in NDC (using specialization constants)
    vec2 tileSize = vec2(float(WORKGROUP_SIZE_X), float(WORKGROUP_SIZE_Y));
    vec2 minNdc = (vec2(tileId) * tileSize) / vec2(pc.screenSize) * 2.0 - 1.0;
    vec2 maxNdc = (vec2(tileId + 1) * tileSize) / vec2(pc.screenSize) * 2.0 - 1.0;

    // View-space corners at near plane
    vec3 corners[4];
    corners[0] = screenToView(vec2(minNdc.x, minNdc.y) * 0.5 + 0.5, 0.0);
    corners[1] = screenToView(vec2(maxNdc.x, minNdc.y) * 0.5 + 0.5, 0.0);
    corners[2] = screenToView(vec2(maxNdc.x, maxNdc.y) * 0.5 + 0.5, 0.0);
    corners[3] = screenToView(vec2(minNdc.x, maxNdc.y) * 0.5 + 0.5, 0.0);

    // Frustum planes (in view space)
    // Left, Right, Bottom, Top, Near, Far
    planes[0] = vec4(normalize(cross(corners[3], corners[0])), 0.0); // Left
    planes[1] = vec4(normalize(cross(corners[1], corners[2])), 0.0); // Right
    planes[2] = vec4(normalize(cross(corners[0], corners[1])), 0.0); // Bottom
    planes[3] = vec4(normalize(cross(corners[2], corners[3])), 0.0); // Top
    planes[4] = vec4(0.0, 0.0, -1.0, -minDepth); // Near
    planes[5] = vec4(0.0, 0.0, 1.0, maxDepth);   // Far
}

// Test if a sphere (light) intersects the frustum
bool sphereInFrustum(vec3 center, float radius, vec4 planes[6]) {
    for (int i = 0; i < 6; i++) {
        float dist = dot(planes[i].xyz, center) + planes[i].w;
        if (dist < -radius) {
            return false;
        }
    }
    return true;
}

void main() {
    uvec2 tileId = gl_WorkGroupID.xy;
    uvec2 localId = gl_LocalInvocationID.xy;
    uint localIndex = localId.y * WORKGROUP_SIZE_X + localId.x;
    
    // Calculate tile index for output buffer
    uvec2 numTiles = (pc.screenSize + uvec2(WORKGROUP_SIZE_X, WORKGROUP_SIZE_Y) - 1) / uvec2(WORKGROUP_SIZE_X, WORKGROUP_SIZE_Y);
    uint tileIndex = tileId.y * numTiles.x + tileId.x;
    uint tileOffset = tileIndex * (MAX_LIGHTS_PER_TILE + 1);
    
    // Initialize shared memory
    if (localIndex == 0) {
        sharedLightCount = 0;
        sharedMinDepthInt = 0xFFFFFFFFu; // Max uint for min comparison
        sharedMaxDepthInt = 0u;           // 0 for max comparison
    }
    barrier();
    
    // Sample depth buffer to find min/max depth for this tile
    uvec2 pixelCoord = tileId * uvec2(WORKGROUP_SIZE_X, WORKGROUP_SIZE_Y) + localId;
    if (pixelCoord.x < pc.screenSize.x && pixelCoord.y < pc.screenSize.y) {
        vec2 uv = (vec2(pixelCoord) + 0.5) / vec2(pc.screenSize);
        float depth = texture(depthBuffer, uv).r;
        uint depthInt = floatBitsToUint(depth);
        
        atomicMin(sharedMinDepthInt, depthInt);
        atomicMax(sharedMaxDepthInt, depthInt);
    }
    barrier();
    
    // Convert back to float
    float minDepth = uintBitsToFloat(sharedMinDepthInt);
    float maxDepth = uintBitsToFloat(sharedMaxDepthInt);
    
    // OPTIMIZATION: Only thread 0 computes frustum, others read from shared memory
    // This reduces frustum calculations from 256 per tile to 1 per tile
    if (localIndex == 0) {
        vec4 frustumPlanes[6];
        getTileFrustum(tileId, minDepth, maxDepth, frustumPlanes);
        // Copy to shared memory
        for (int i = 0; i < 6; i++) {
            sharedFrustumPlanes[i] = frustumPlanes[i];
        }
    }
    barrier();
    
    // OPTIMIZATION (Issue 20 partial): Early-exit if no lights to cull
    // Full adaptive workgroup sizing would require specialization constants
    // TODO: Host-side should skip dispatch entirely if lightCount == 0
    if (pc.lightCount == 0) {
        if (localIndex == 0) {
            tileData[tileOffset] = 0;
        }
        return;
    }
    
    // Each thread tests some lights
    // Use actual workgroup size from specialization constants
    uint workgroupSize = WORKGROUP_SIZE_X * WORKGROUP_SIZE_Y;
    uint lightsPerThread = (pc.lightCount + workgroupSize - 1) / workgroupSize;
    uint lightStart = localIndex * lightsPerThread;
    uint lightEnd = min(lightStart + lightsPerThread, pc.lightCount);
    
    for (uint i = lightStart; i < lightEnd; i++) {
        Light light = lights[i];
        
        // Skip disabled lights
        if (light.params.w < 0.5) continue;
        
        // Transform light to view space
        vec3 lightPosView = (camera.view * vec4(light.position.xyz, 1.0)).xyz;
        float lightRadius = light.position.w;
        
        // Test against frustum (using shared memory)
        if (sphereInFrustum(lightPosView, lightRadius, sharedFrustumPlanes)) {
            uint slot = atomicAdd(sharedLightCount, 1);
            if (slot < MAX_LIGHTS_PER_TILE) {
                sharedLightIndices[slot] = i;
            }
        }
    }
    barrier();
    
    // Write results to global buffer (only thread 0)
    if (localIndex == 0) {
        uint count = min(sharedLightCount, MAX_LIGHTS_PER_TILE);
        tileData[tileOffset] = count;
        
        for (uint i = 0; i < count; i++) {
            tileData[tileOffset + 1 + i] = sharedLightIndices[i];
        }
    }
}
