struct VertexInput {
    @location(0) position: vec3<f32>,
    @location(1) color: vec3<f32>,
    @location(2) normal: vec3<f32>,
};

const CLUSTER_COUNT_X: u32 = 16u;
const CLUSTER_COUNT_Y: u32 = 9u;
const CLUSTER_COUNT_Z: u32 = 24u;

const NEAR_PLANE: f32 = 0.1;
const FAR_PLANE: f32 = 1000.0;

struct InstanceInput {
    @location(5) model_row0: vec4<f32>,
    @location(6) model_row1: vec4<f32>,
    @location(7) model_row2: vec4<f32>,
    @location(8) model_row3: vec4<f32>,
    @location(9) color: vec3<f32>,
    @location(10) flags: u32,
};

struct VSOutput {
    @builtin(position) position: vec4<f32>,
    @location(0) frag_color: vec3<f32>,
    @location(1) world_pos: vec3<f32>,
    @location(2) normal: vec3<f32>,
    @location(3) flags: u32,
};

struct LightCount {
    count: u32,
};

@group(1) @binding(1)
var<uniform> light_count: LightCount;

struct Globals {
    view_proj: mat4x4<f32>,
    resolution: vec2<f32>,
}

@group(0) @binding(0)
var<uniform> globals: Globals;

struct GpuLight {
    position: vec3<f32>,
    light_type: u32,
    color: vec3<f32>,
    intensity: f32,
    direction: vec3<f32>,
    range: f32,
    inner_cutoff: f32,
    outer_cutoff: f32,
};

@group(1) @binding(0)
var<storage, read> all_lights: array<GpuLight>;

@group(2) @binding(0)
var<storage, read> cluster_light_indices: array<u32>;
@group(2) @binding(1)
var<storage, read> cluster_offsets: array<vec2<u32>>;

@vertex
fn vs_main(vertex: VertexInput, instance: InstanceInput) -> VSOutput {
    var out: VSOutput;

    let model = mat4x4<f32>(
        instance.model_row0,
        instance.model_row1,
        instance.model_row2,
        instance.model_row3
    );

    let world_position = (model * vec4<f32>(vertex.position, 1.0)).xyz;
    let normal_matrix = mat3x3<f32>(
        instance.model_row0.xyz,
        instance.model_row1.xyz,
        instance.model_row2.xyz
    );

    out.position = globals.view_proj * vec4<f32>(world_position, 1.0);
    out.frag_color = instance.color * vertex.color;
    out.world_pos = world_position;
    out.normal = normalize(normal_matrix * vertex.normal);
    out.flags = instance.flags;

    return out;
}

fn compute_cluster_id(frag_coord: vec4<f32>, view_pos_z: f32, screen_size: vec2<f32>) -> u32 {
    let x_frac = frag_coord.x / screen_size.x;
    let y_frac = frag_coord.y / screen_size.y;

    let x = clamp(u32(x_frac * f32(CLUSTER_COUNT_X)), 0u, CLUSTER_COUNT_X - 1u);
    let y = clamp(u32(y_frac * f32(CLUSTER_COUNT_Y)), 0u, CLUSTER_COUNT_Y - 1u);

    // Z: logarithmic depth
    let depth = -view_pos_z; // view-space z is negative
    let depth_clamped = clamp(depth, NEAR_PLANE, FAR_PLANE);
    let log_depth = log2(depth_clamped);
    let z = clamp(u32((log_depth / log2(FAR_PLANE / NEAR_PLANE)) * f32(CLUSTER_COUNT_Z)), 0u, CLUSTER_COUNT_Z - 1u);

    return x + y * CLUSTER_COUNT_X + z * CLUSTER_COUNT_X * CLUSTER_COUNT_Y;
}

fn is_nan_f32(x: f32) -> bool {
    return x != x;
}

fn is_nan_vec3(v: vec3<f32>) -> bool {
    return any(vec3<bool>(v != v));
}

@fragment
fn fs_main(input: VSOutput) -> @location(0) vec4<f32> {
    var lighting: vec3<f32> = vec3<f32>(0.0);
    let always_lit = (input.flags & 0x1u) != 0u;

    let cluster_id = compute_cluster_id(input.position, input.world_pos.z, globals.resolution);
    let offset_info = cluster_offsets[cluster_id];
    let offset = offset_info.x;
    let count = offset_info.y;

    for (var i = 0u; i < count; i = i + 1u) {
        let light_index = cluster_light_indices[offset + i];
        let light = all_lights[light_index];
        var light_contrib: vec3<f32> = vec3<f32>(0.0);

        let light_dir = normalize(light.position - input.world_pos);
        let diff = max(dot(input.normal, light_dir), 0.0);

        switch (light.light_type) {
            case 0u: { // Directional
                light_contrib = light.color * light.intensity * diff;
            }
            case 1u: { // Point
                let dist = distance(light.position, input.world_pos);
                if (dist < light.range) {
                    let attenuation = 1.0 / (dist * dist);
                    light_contrib = light.color * light.intensity * diff * attenuation;
                }
            }
            case 2u: { // Spot
                let spot_dir = normalize(-light.direction);
                let angle = dot(spot_dir, light_dir);
                if (angle > light.outer_cutoff) {
                    let intensity = clamp((angle - light.outer_cutoff) / (light.inner_cutoff - light.outer_cutoff), 0.0, 1.0);
                    let dist = distance(light.position, input.world_pos);
                    let attenuation = 1.0 / (dist * dist);
                    light_contrib = light.color * light.intensity * diff * attenuation * intensity;
                }
            }
            default: {}
        }

        if (!always_lit) {
            lighting += light_contrib;
        }
    }

    if (always_lit) {
        lighting = vec3<f32>(1.0, 1.0, 1.0) * 2.0;
    }

    return vec4<f32>(input.frag_color * lighting, 1.0);
}