// BVLE Voxels - RT Shadow + AO Compute Shader (Phase 6.2 + 6.3)
// Per-pixel: traces 1 shadow ray toward sun + N hemisphere rays for AO.
// Modulates voxelRT_ in-place via RWTexture2D.

#include "voxelCommon.hlsli"

// SRV bindings
Texture2D<float>  depthTexture  : register(t0);  // voxelDepth_ (D32_FLOAT as R32_FLOAT SRV)
Texture2D<float4> normalTexture : register(t1);  // voxelNormalRT_ (R16G16B16A16_SNORM)
RaytracingAccelerationStructure tlas : register(t2);  // TLAS with blocky + smooth instances

// UAV outputs
RWTexture2D<float4> colorOutput : register(u0);  // voxelRT_ (shadow applied in-place)
RWTexture2D<float>  aoOutput    : register(u1);  // raw AO factor (blurred separately)

// Push constants
struct ShadowPush {
    uint width;
    uint height;
    float normalBias;
    float shadowMaxDist;
    uint debugMode;     // 0=normal, 1=debug shadows, 2=debug AO
    float aoRadius;     // max distance for AO rays (e.g. 8.0 voxels)
    uint aoRayCount;    // number of hemisphere rays (e.g. 6)
    float aoStrength;   // how dark full occlusion is (e.g. 0.35 = 65% darkening)
    uint pad[4];
};
[[vk::push_constant]] ConstantBuffer<ShadowPush> push : register(b999);

// ── Hash-based pseudo-random for AO ray directions ──────────────
// Golden ratio hash: deterministic, no texture lookup, good distribution
float hashF(uint x) {
    x ^= x >> 16;
    x *= 0x45d9f3bu;
    x ^= x >> 16;
    return float(x & 0xFFFFFF) / float(0xFFFFFF);
}

uint hashU(uint a, uint b) {
    a ^= b * 0x9E3779B9u;
    a ^= a >> 16;
    a *= 0x45d9f3bu;
    return a;
}

// Build orthonormal basis from normal (Frisvad's method, robust for all N)
void buildBasis(float3 N, out float3 T, out float3 B) {
    if (N.z < -0.9999) {
        T = float3(0, -1, 0);
        B = float3(-1, 0, 0);
    } else {
        float a = 1.0 / (1.0 + N.z);
        float b = -N.x * N.y * a;
        T = float3(1.0 - N.x * N.x * a, b, -N.x);
        B = float3(b, 1.0 - N.y * N.y * a, -N.y);
    }
}

// Cosine-weighted hemisphere sample (probability ∝ cos(θ))
float3 cosineSampleHemisphere(float u1, float u2, float3 N, float3 T, float3 B) {
    float r = sqrt(u1);
    float phi = 6.28318530718 * u2;
    float x = r * cos(phi);
    float y = r * sin(phi);
    float z = sqrt(max(0.0, 1.0 - u1));
    return normalize(x * T + y * B + z * N);
}

[RootSignature(VOXEL_ROOTSIG)]
[numthreads(8, 8, 1)]
void main(uint3 DTid : SV_DispatchThreadID) {
    if (DTid.x >= push.width || DTid.y >= push.height) return;

    float depth = depthTexture[DTid.xy];
    if (depth == 0.0) {
        if (push.debugMode > 0) colorOutput[DTid.xy] = float4(0.1, 0.1, 0.1, 1);
        return;
    }

    // Reconstruct world position from depth via inverse VP
    float2 uv = (float2(DTid.xy) + 0.5) / float2(push.width, push.height);
    float2 ndc = float2(uv.x * 2.0 - 1.0, (1.0 - uv.y) * 2.0 - 1.0);
    float4 clipPos = float4(ndc, depth, 1.0);
    float4 worldPos4 = mul(inverseViewProjection, clipPos);
    float3 worldPos = worldPos4.xyz / worldPos4.w;

    float3 N = normalTexture[DTid.xy].xyz;
    float3 origin = worldPos + N * push.normalBias;

    // ── Shadow ray toward sun ──────────────────────────────────
    float3 L = normalize(-sunDirection.xyz);
    float NdotL = dot(N, L);

    float shadowFactor = 1.0;
    if (NdotL <= 0.0) {
        shadowFactor = 0.3; // back-facing = fully in shadow
    } else {
        RayDesc ray;
        ray.Origin = origin;
        ray.Direction = L;
        ray.TMin = 0.01;
        ray.TMax = push.shadowMaxDist;

        RayQuery<RAY_FLAG_SKIP_PROCEDURAL_PRIMITIVES | RAY_FLAG_ACCEPT_FIRST_HIT_AND_END_SEARCH> q;
        q.TraceRayInline(tlas, 0, 0xFF, ray);
        [loop] while (q.Proceed()) {}

        if (q.CommittedStatus() == COMMITTED_TRIANGLE_HIT) {
            shadowFactor = 0.3;
        }
    }

    // ── AO: hemisphere rays ────────────────────────────────────
    float aoFactor = 1.0;
    uint rayCount = push.aoRayCount;
    if (rayCount > 0) {
        float3 T, B;
        buildBasis(N, T, B);

        // Fully world-space seed: solid voxel coord + tangent-plane frac position
        // vc: offset by -N*0.5 → inside the solid voxel (stable per face)
        // Sub-voxel: only use the 2 tangent axes (T,B), NOT the normal axis.
        // The normal axis sits at an integer boundary (e.g. face +Y → y=5.0000)
        // where frac() oscillates between ~0 and ~1 due to depth precision → flicker.
        // Tangent axes vary smoothly across the face → always stable.
        int3 vc = int3(floor(worldPos - N * 0.5));
        float tFrac = frac(dot(worldPos, T));
        float bFrac = frac(dot(worldPos, B));
        uint st = uint(tFrac * 256.0);
        uint sb = uint(bFrac * 256.0);
        uint baseSeed = hashU(hashU((uint)(vc.x + 32768), (uint)(vc.y + 32768)), (uint)(vc.z + 32768));
        uint pixelSeed = hashU(baseSeed, hashU(st, sb));
        float totalOcclusion = 0.0;

        [loop]
        for (uint i = 0; i < rayCount; i++) {
            // Per-ray random: hash(pixelSeed, rayIndex)
            uint seed = hashU(pixelSeed, i);
            float u1 = hashF(seed);
            float u2 = hashF(seed ^ 0xA5A5A5A5u);

            float3 dir = cosineSampleHemisphere(u1, u2, N, T, B);

            RayDesc aoRay;
            aoRay.Origin = origin;
            aoRay.Direction = dir;
            aoRay.TMin = 0.05; // larger TMin for AO to avoid edge self-intersection
            aoRay.TMax = push.aoRadius;

            RayQuery<RAY_FLAG_SKIP_PROCEDURAL_PRIMITIVES | RAY_FLAG_ACCEPT_FIRST_HIT_AND_END_SEARCH> aoQ;
            aoQ.TraceRayInline(tlas, 0, 0xFF, aoRay);
            [loop] while (aoQ.Proceed()) {}

            if (aoQ.CommittedStatus() == COMMITTED_TRIANGLE_HIT) {
                // Distance-weighted: close hits = strong occlusion, far hits = weak
                float hitT = aoQ.CommittedRayT();
                float falloff = 1.0 - saturate(hitT / push.aoRadius);
                totalOcclusion += falloff * falloff; // quadratic for natural falloff
            }
        }

        float occlusionRatio = totalOcclusion / float(rayCount);
        aoFactor = 1.0 - occlusionRatio * push.aoStrength;
    }

    // ── Write AO to separate buffer (will be blurred), apply shadow in-place ──
    aoOutput[DTid.xy] = aoFactor;

    if (push.debugMode == 1) {
        // Debug shadows: red=shadow, green=lit, blue=backface
        if (NdotL <= 0.0)
            colorOutput[DTid.xy] = float4(0, 0, 0.5, 1);
        else if (shadowFactor < 1.0)
            colorOutput[DTid.xy] = float4(1, 0, 0, 1);
        else
            colorOutput[DTid.xy] = float4(0, 1, 0, 1);
    } else if (push.debugMode == 2) {
        // Debug AO: raw AO written to aoOutput, will be visualized after blur
        // Write white to color so blur apply pass shows AO only
        colorOutput[DTid.xy] = float4(1, 1, 1, 1);
    } else {
        // Apply shadow only — AO applied after blur in a separate pass
        float4 color = colorOutput[DTid.xy];
        color.rgb *= shadowFactor;
        colorOutput[DTid.xy] = color;
    }
}