#include "sh/common.hlsl"

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

#define ROOTSIG \
    "RootConstants(num32BitConstants = 7, b0), " \
    "DescriptorTable(SRV(t0, space = 0, numDescriptors = unbounded, flags = DESCRIPTORS_VOLATILE)), " \
    "DescriptorTable(SRV(t0, space = 1, numDescriptors = unbounded, flags = DESCRIPTORS_VOLATILE)), " \
    "DescriptorTable(UAV(u0, space = 2, 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_shade_constants> g_constants : register(b0);
Texture2D<float4> g_gbuff_textures[] : register(t0, space0);
Texture2D<uint2> g_emittance_flood_textures[] : register(t0, space1);
RWTexture2D<float4> g_write_textures[]: register(u0, space2);

SamplerState g_sampler : register(s0);

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

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

#define SAMPLES 4
#define MARCHES 16

INLINE float4 get_light_in_dir(uint2 pos, float2 dir)
{
    float4 result = 0;
    uint2 at = pos;
    for (uint i = 0; i < MARCHES; ++i) {
        uint2 flood = g_emittance_flood_textures[g_constants.emittance_flood_tex_urid][at];
        int2 dist_vec = (int2)at - (int2)flood;
        float dist = length(dist_vec);
        if (dist <= 2) {
            result = g_gbuff_textures[g_constants.emittance_tex_urid][flood];
            break;
        } else {
            at += dir * dist;
        }
    }
    return result;
}

float rand_float_from_float2(float2 pos) {
    return frac(sin(dot(pos.xy, float2(12.9898,78.233))) * 43758.5453123);
}

INLINE float4 get_light_at_pos(uint2 pos)
{
    float4 result = 0;
    for (int i = 0; i < SAMPLES; ++i) {
        float angle = TAU * (((float)i + rand_float_from_float2(pos + (float)i + sin(g_constants.time))) / ((float)SAMPLES - 1));
        float2 dir = float2(sin(angle), cos(angle));
        float4 light_in_dir = get_light_in_dir(pos, dir);
        result += light_in_dir;
    }
    result /= SAMPLES;

    return result;
}

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

    float4 final_color = albedo * lighting;

    g_write_textures[g_constants.write_tex_urid][id] = final_color;
}