// 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 "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;

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: proxy box for each object transformed to local space.
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 lp       = (obj.inv_model * vec4f(p, 1.0)).xyz;
        let obj_type = u32(obj.params.x);
        var d: f32;
        switch obj_type {
            case 1u: { d = sdSphere(lp, 1.0); }                              // SPHERE
            case 2u: { d = sdPlane(lp, vec3f(0.0, 1.0, 0.0), obj.params.y); } // PLANE
            case 3u: { d = sdTorus(lp, vec2f(0.8, 0.2)); }                   // TORUS
            default: { d = sdBox(lp, vec3f(1.0)); }                          // CUBE (0) + fallback
        }
        if (d < res.distance) {
            res.distance  = d;
            res.object_id = f32(i + 1u);
        }
    }
    return res;
}

// ---- 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;

    // Surface normal estimated from SDF gradient.
    let nor      = normalWithID(world);
    let bias_pos = world + nor * 0.02;

    // 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  = shadowWithStoredDistance(bias_pos, ld, MAX_RAY_LENGTH);
        shadow_val = min(shadow_val, s);
    }

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