#include "sh/common.hlsl"

/* ========================== *
 * Root signature
 * ========================== */

#define ROOTSIG \
    "RootConstants(num32BitConstants = 6, b0), " \
    "DescriptorTable(SRV(t0, space = 0, numDescriptors = unbounded, flags = DESCRIPTORS_VOLATILE)), " \
    "DescriptorTable(UAV(u0, space = 1, numDescriptors = unbounded, flags = DESCRIPTORS_VOLATILE)), " \
    \
    \
    "StaticSampler(s0, " \
                  "filter = FILTER_MIN_MAG_MIP_POINT, " \
                  "addressU = TEXTURE_ADDRESS_CLAMP, " \
                  "addressV = TEXTURE_ADDRESS_CLAMP, " \
                  "addressW = TEXTURE_ADDRESS_CLAMP, " \
                  "maxAnisotropy = 1)"

ConstantBuffer<struct sh_flood_constants> g_constants : register(b0);
Texture2D<float4> g_emittance_textures[] : register(t0, space0);
RWTexture2D<int2> g_flood_textures[]: register(u0, space1);

SamplerState g_sampler : register(s0);

/* ========================== *
 * Compute shader
 * ========================== */

struct cs_input {
    DECLS(uint3, SV_DispatchThreadID);
};

[numthreads(8, 8, 1)]
SH_ENTRY(ROOTSIG) void cs(struct cs_input input)
{
    uint2 id = input.SV_DispatchThreadID.xy;
    uint2 tex_size = uint2(g_constants.tex_width, g_constants.tex_height);
    if (id.x >= tex_size.x || id.y >= tex_size.y) {
        return;  /* Overflow */
    }

    int step_len = g_constants.step_len;
    if (step_len == -1) {
        /* Seed */
        float4 emittance = g_emittance_textures[g_constants.emittance_tex_urid][id];
        int2 seed = int2(-1, -1);
        if (emittance.a > 0) {
            seed = id.xy;
        }
        g_flood_textures[g_constants.flood_read_tex_urid][id] = seed;
        g_flood_textures[g_constants.flood_write_tex_urid][id] = seed;
    } else {
        /* Flood */
        int2 seed = g_flood_textures[g_constants.flood_read_tex_urid][id];
        if (seed.x >= 0 && seed.y >= 0) {
            int2 flood_coords[2] = {
                int2((int)id.x - step_len,  (int)id.y           ),  /* cl */
                int2((int)id.x + step_len,  (int)id.y           ),  /* cr */
            };
            for (int i = 0; i < 2; ++i) {
                int2 coord = flood_coords[i];
                if (coord.x >= 0 && coord.x < (int)tex_size.x && coord.y >= 0 && coord.y < (int)tex_size.y) {
                    int2 old_flood = g_flood_textures[g_constants.flood_read_tex_urid][coord];
                    // if (old_flood.x < 0 || old_flood.y < 0) {
                        g_flood_textures[g_constants.flood_write_tex_urid][coord] = seed;
                    // }
                }
            }
        }
        // if (seed.x >= 0 && seed.y >= 0) {
        //     int2 flood_coords[8] = {
        //         int2((int)id.x - step_len,  (int)id.y - step_len),  /* tl */
        //         int2((int)id.x           ,  (int)id.y - step_len),  /* tc */
        //         int2((int)id.x + step_len,  (int)id.y - step_len),  /* tr */
        //         int2((int)id.x - step_len,  (int)id.y           ),  /* cl */
        //         int2((int)id.x + step_len,  (int)id.y           ),  /* cr */
        //         int2((int)id.x - step_len,  (int)id.y + step_len),  /* bl */
        //         int2((int)id.x           ,  (int)id.y + step_len),  /* bc */
        //         int2((int)id.x + step_len,  (int)id.y + step_len)   /* br */
        //     };
        //     for (int i = 0; i < 8; ++i) {
        //         int2 coord = flood_coords[i];
        //         if (coord.x >= 0 && coord.x < (int)tex_size.x && coord.y >= 0 && coord.y < (int)tex_size.y) {
        //             int2 old_flood = g_flood_textures[g_constants.flood_read_tex_urid][coord];
        //             if (old_flood.x < 0 || old_flood.y < 0) {
        //                 g_flood_textures[g_constants.flood_write_tex_urid][coord] = seed;
        //             }
        //         }
        //     }
        // }
    }

}