// G-buffer shadow raymarching shader for CNN v3
// Pass 2: Reads depth from Pass 1, marches shadow rays toward lights,
// outputs shadow factor (1.0=lit, 0.0=shadow) to RGBA8Unorm render target (.r).

#include "common_uniforms"
#include "camera_common"
#include "math/sdf_shapes"
#include "math/normal"
#include "render/raymarching_id"

@group(0) @binding(0) var<uniform>         globals:     GlobalUniforms;
@group(0) @binding(1) var<storage, read>   object_data: ObjectsBuffer;
@group(0) @binding(2) var                  depth_tex:   texture_depth_2d;
@group(0) @binding(4) var                  normal_mat_tex: texture_2d<f32>;

struct GBufLight {
    direction: vec4f,  // xyz = toward light (world space, normalized)
    color:     vec4f,  // rgb = color, a = intensity
}
struct GBufLightsUniforms {
    lights: array<GBufLight, 2>,
    params: vec4f,  // x = num_lights
}
@group(0) @binding(3) var<uniform> lights: GBufLightsUniforms;

// ---- SDF scene (proxy box per object in local space) ----

// Stub required by render/raymarching (shadow() / rayMarch() call df()).
fn df(p: vec3f) -> f32 { return MAX_RAY_LENGTH; }

// SDF of the full scene.
// Sphere: direct world-space formula (exact, no matrix multiply).
// Box/Torus/Plane: local-space transform + uniform-scale correction.
fn dfWithID(p: vec3f) -> RayMarchResult {
    var res: RayMarchResult;
    res.distance     = MAX_RAY_LENGTH;
    res.distance_max = MAX_RAY_LENGTH;
    res.object_id    = 0.0;

    let n = u32(globals.params.x);
    for (var i = 0u; i < n; i++) {
        let obj      = object_data.objects[i];
        let obj_type = u32(obj.params.x);
        var d: f32;
        switch obj_type {
            case 1u: {
                // SPHERE: direct world-space SDF — avoids matrix multiply, exact.
                let c = obj.model[3].xyz;
                let r = length(obj.model[0].xyz);
                d = length(p - c) - r;
            }
            case 2u: {
                // PLANE
                let lp = (obj.inv_model * vec4f(p, 1.0)).xyz;
                d = sdPlane(lp, vec3f(0.0, 1.0, 0.0), obj.params.y);
            }
            case 3u: {
                // TORUS
                let lp    = (obj.inv_model * vec4f(p, 1.0)).xyz;
                let scale = length(obj.model[0].xyz);
                d = sdTorus(lp, vec2f(0.8, 0.2)) * scale;
            }
            default: {
                // CUBE (0) + fallback — uniform scale assumed.
                let lp    = (obj.inv_model * vec4f(p, 1.0)).xyz;
                let scale = length(obj.model[0].xyz);
                d = sdBox(lp, vec3f(1.0)) * scale;
            }
        }
        if (d < res.distance) {
            res.distance  = d;
            res.object_id = f32(i + 1u);
        }
    }
    return res;
}

// Soft shadow march (IQ formula). Returns 1=lit, 0=shadow.
// No dmin/dmax bounds: in open space d grows large so 8*d/t >> 1, res stays 1 naturally.
fn soft_shadow(ro: vec3f, rd: vec3f) -> f32 {
    var t   = 0.001;
    var res = 1.0;
    for (var i = 0; i < 64; i++) {
        let d = dfWithID(ro + rd * t).distance;
        if (d < 0.0005) { return 0.0; }
        res = min(res, 8.0 * d / t);
        t  += d;
    }
    return clamp(res, 0.0, 1.0);
}

// ---- Vertex: fullscreen triangle ----

@vertex
fn vs_main(@builtin(vertex_index) vid: u32) -> @builtin(position) vec4f {
    let x = f32((vid & 1u) << 2u) - 1.0;
    let y = f32((vid & 2u) << 1u) - 1.0;
    return vec4f(x, y, 0.0, 1.0);
}

// ---- Fragment: shadow factor per pixel ----

@fragment
fn fs_main(@builtin(position) pos: vec4f) -> @location(0) vec4f {
    let depth = textureLoad(depth_tex, vec2i(pos.xy), 0);

    // Sky / background: fully lit.
    if (depth >= 1.0) {
        return vec4f(1.0);
    }

    // Reconstruct world-space position from NDC + depth.
    let res = globals.resolution;
    let ndc = vec2f(
        (pos.x / res.x) * 2.0 - 1.0,
        1.0 - (pos.y / res.y) * 2.0
    );
    let clip  = globals.inv_view_proj * vec4f(ndc, depth, 1.0);
    let world = clip.xyz / clip.w;

    // Use rasterized surface normal for bias — correct for sphere impostors.
    let nm       = textureLoad(normal_mat_tex, vec2i(pos.xy), 0);
    let nor      = oct_decode_unorm(nm.rg);
    let bias_pos = world + nor * 0.05;

    // March shadow rays toward each light; take the darkest value.
    var shadow_val  = 1.0;
    let num_lights  = u32(lights.params.x);
    for (var i = 0u; i < num_lights; i++) {
        let ld = lights.lights[i].direction.xyz;
        let s  = soft_shadow(bias_pos, ld);
        shadow_val = min(shadow_val, s);
    }

    return vec4f(shadow_val, shadow_val, shadow_val, 1.0);
}