struct PushConstants
{
    float4x4 view_proj;
    float3   cam_pos;
    uint32_t frame_timestamp;
    uint32_t width;
    uint32_t height;

    uint32_t chunk_width;
    uint32_t chunk_height;
    uint32_t chunk_alt;
}

public struct VertexOutput
{
    public float4 position : SV_Position;

    [vk::location(0)]
    public float3 world_position;

    [vk::location(1)]
    public nointerpolation uint32_t structure_id;

    [vk::location(2)]
    public float3 cam_position;

    [vk::location(3)]
    public float3 chunk_position;
}

[[vk::push_constant]]
uniform PushConstants constants;

[shader("vertex")]
VertexOutput chunk(
    uint                     index: SV_VulkanVertexID,
    [vk::location(0)] float3 chunk_position,
    [vk::location(1)] uint   id)
{
    let cube_vertices : float3[8] =
                            float3[](
                                float3(0., 0., 0.),
                                float3(0., 0., 1.),
                                float3(1., 0., 1.),
                                float3(1., 0., 0.),

                                float3(0., 1., 0.),
                                float3(0., 1., 1.),
                                float3(1., 1., 1.),
                                float3(1., 1., 0.), );
    // clang-format off
    let cube_faces: int[24] = int[](
        // Bottom face
        1, 0, 2, 3,

        // Top face
        4, 5, 7, 6,

        // Side faces
        0, 1, 4, 5,
        1, 2, 5, 6,
        2, 3, 6, 7,
        3, 0, 7, 4,
    );

    let quad_index = index / (3 * 2);
    let triangle_index = index % (3 * 2);
    let triangle_map: int[6] = int[](
        0, 1, 2, 1, 3, 2
    );


    let vertex = cube_vertices[cube_faces[quad_index * 4 + triangle_map[triangle_index]]];
    let output_vertex = mul(constants.view_proj, float4(vertex + chunk_position, 1.0f));

    VertexOutput vertex_output;
    vertex_output.position = output_vertex;
    vertex_output.world_position = vertex + chunk_position;
    vertex_output.structure_id = id;
    vertex_output.cam_position = constants.cam_pos;
    vertex_output.chunk_position = chunk_position;

    return vertex_output;
}

struct StructurePointer
{
    uint32_t value;
    bool subdivided()
    {
        return (this.value & 0x80000000) != 0;
    }

    bool pointer_valid()
    {
        return (this.value & 0x40000000) != 0;
    }

    uint32_t pointer()
    {
        return this.value & 0x3FFFFFFF;
    }
}

struct ByteColor
{
    uint32_t byte_color;
    
    property uint32_t byte_r {
        get {return byte_color & 0xFF;}
    }

    property uint32_t byte_g {
        get {return (byte_color >> 8) & 0xFF;}
    }

    property uint32_t byte_b {
        get {return (byte_color >> 16) & 0xFF;}
    }

    property uint32_t byte_a {
        get {return byte_color >> 24;}
    }

    property float4 float_color {
        get {return float4(
                float(byte_r) / 255.,
                float(byte_g) / 255.,
                float(byte_b) / 255.,
                float(byte_a) / 255.
        ); }
    }
}

struct StructurePoolElement
{
    uint32_t occupancy_low;
    uint32_t occupancy_high;
    StructurePointer pointers[64];
}

struct RequestBufferElement
{
    Atomic<uint32_t> requests[64];
}

struct ColorPoolElement
{
    ByteColor colors[64];
}

struct LocationPoolElement
{
    uint32_t structure_id;
    uint32_t structure_locator;
}

[[vk::binding(0, 0)]] RWStructuredBuffer<StructurePoolElement> structure_pool;
[[vk::binding(1, 0)]] RWStructuredBuffer<ColorPoolElement> color_pool;
[[vk::binding(2, 0)]] RWStructuredBuffer<LocationPoolElement> location_pool;
[[vk::binding(3, 0)]] RWStructuredBuffer<RequestBufferElement> request_buffer;
[[vk::binding(4, 0)]] RWStructuredBuffer<uint32_t> usage_buffer;
[[vk::binding(5, 0)]] RWStructuredBuffer<StructurePointer> structure_table_pointer;
[[vk::binding(6, 0)]] RWStructuredBuffer<Atomic<uint32_t>> structure_table_request_buffer;

uint32_t3 get_children_pos(float3 position, uint32_t scale_exp)
{
    return (asuint(position) >> scale_exp) & 3;
}

uint32_t get_children_index(float3 position, uint32_t scale_exp)
{
    // Get mantissa bits for this scale exp an retain bits for the specific children
    uint32_t3 cell_position = (asuint(position) >> scale_exp) & 3;
    return cell_position.x + cell_position.y * 4 + cell_position.z * 4 * 4;
}

uint64_t get_child_mask(uint32_t low, uint32_t high)
{
    return ((uint64_t)high << 32) | (uint64_t)low;
}

float3 floor_scale(float3 position, uint32_t scale_exp)
{
    uint32_t mask = ~0u << scale_exp;
    return asfloat(asuint(position) & mask);
}

struct HitInformation
{
    bool hit;
    float3 hit_pos;
    float4 color;
}

HitInformation ray_march(float3 ray_direction, float3 ray_origin, uint32_t root_id, float dist_offset)
{
    float fov_deg = 100. / 1920.;
    float fov_rad = (float.getPi() * fov_deg) / 180.;
    float cone_factor = tan(fov_rad / 2.) * 2; // Horizontal size of pixel

    let st_pointer = structure_table_pointer[root_id];
    if(!st_pointer.subdivided())
    {
        var hit: HitInformation;
        hit.hit = false;
        return hit;
    }

    if(!st_pointer.pointer_valid())
    {
        // Record request
        structure_table_request_buffer[root_id].add(1);
        var hit: HitInformation;
        hit.hit = false;
        return hit;
    }

    ray_origin += float3(1.);
    float3 pos = ray_origin;
    pos = clamp(pos, float(1.), asfloat(0x3fffffff));

    uint32_t scale_exp = 23 - 2;
    uint32_t node_stack[5] = 
    {
        0
    };
    uint32_t current_node_index = structure_table_pointer[root_id].pointer();

    node_stack[10 - scale_exp / 2] = current_node_index;

    uint32_t child_index = get_children_index(pos, scale_exp);
    StructurePointer current_node = structure_pool[current_node_index].pointers[child_index];

    usage_buffer[current_node_index] = constants.frame_timestamp;


    for(uint32_t iter = 0; iter < 500; iter ++)
    {
        //scale_exp = 23 - 2;
        //current_node_index = structure_table_pointer[root_id].pointer();

        child_index = get_children_index(pos, scale_exp);
        current_node = structure_pool[current_node_index].pointers[child_index];

        // Scale computations
        //return float4(cone_size * 1000.);
        let cone_size = (length(ray_origin - pos) + dist_offset) * cone_factor;
        let exponent = 
            select(
                cone_size == 0.,
                0,
                23 - (127 - (asuint(cone_size) >> 23))
            );

        while(
            current_node.subdivided() && 
            current_node.pointer_valid() &&
            scale_exp - 2 > exponent
            )
        {
            scale_exp -= 2;
            current_node_index = current_node.pointer();
            node_stack[10 - scale_exp / 2] = current_node_index;
            child_index = get_children_index(pos, scale_exp);
            current_node = structure_pool[current_node_index].pointers[child_index];
            // Write usage
            usage_buffer[current_node_index] = constants.frame_timestamp;
        }

        // Request subdiv
        if(current_node.subdivided() && !current_node.pointer_valid())
        {
            request_buffer[current_node_index].requests[child_index].add(1);
        }

        if(color_pool[current_node_index].colors[child_index].byte_a != 0)
        {
            //hit_pos = pos - float3(1.);
            //return color_pool[current_node_index].colors[child_index].float_color; 
            var hit: HitInformation;
            hit.hit = true;
            hit.hit_pos = pos - float3(1.);
            hit.color = color_pool[current_node_index].colors[child_index].float_color;
            return hit;
        }

        //uint64_t occupancy = get_child_mask(structure_pool[current_node_index].occupancy_low, structure_pool[current_node_index].occupancy_high);
        uint64_t occupancy = get_child_mask(structure_pool[current_node_index].occupancy_low, structure_pool[current_node_index].occupancy_high);
        uint32_t adv_scale_exp = scale_exp;
        if(((occupancy >> (child_index & 0b101010)) & 0x00330033) == 0)
        {
            adv_scale_exp ++;
        }

        // Perform dda
        // Compute correct exponent, and shift it into the exponent part of floatt
        let child_scale : float = asfloat((adv_scale_exp - 23 + 127) << 23);
        let child_pos : float3 = floor_scale(pos, adv_scale_exp);
        let child_far : float3 = child_pos + select(ray_direction > 0., float3(child_scale), float3(0.));
        // Intersection t
        let inter_ts : float3 = (child_far - ray_origin) / ray_direction;
        float inter_t = min(inter_ts.x, min(inter_ts.y, inter_ts.z));
        //return float4(inter_t);

        // Perform dda step
        let neighbor_min : float3 = select(float3(inter_t) == inter_ts, child_pos + copysign(child_scale, ray_direction), child_pos);
        let neighbor_max : float3 = asfloat(asint(neighbor_min) + ((1 << adv_scale_exp) - 1));
        let previous_pos : float3 = pos;
        pos = clamp(ray_origin + ray_direction * inter_t, neighbor_min, neighbor_max);

        /*
        if(any(pos >= 2.) || any(pos < 1.))
        {
            discard;
        }
        */

        // Find most common ancestor
        uint32_t3 diffs = asuint(child_pos) ^ asuint(pos);
        uint32_t diff = (diffs.x | diffs.y | diffs.z);
        
        int32_t common_depth = (1 + (22 - firstbithigh(diff)) / 2) * 2;
        if(common_depth <= 0)
        {
            break;
        }

        scale_exp = 23 - common_depth;
        current_node_index = node_stack[10 - scale_exp / 2];
    }


    var hit: HitInformation;
    hit.hit = false;
    return hit;
}

struct FragmentOutput
{
    float depth : SV_Depth;
    float4 color : SV_Target<0>;
}

//[earlydepthstencil]
[shader("fragment")]
FragmentOutput fragment(VertexOutput vertex_out)
{
    let ray_direction = normalize(vertex_out.world_position - vertex_out.cam_position);
    let intersection_t = box_intersect(vertex_out.cam_position, ray_direction, vertex_out.chunk_position, vertex_out.chunk_position + float3(1.));
    let local_ray_origin = max(intersection_t.x, 0.) * ray_direction + vertex_out.cam_position - vertex_out.chunk_position;
    // Figure out intersection
    let hit = ray_march(ray_direction, local_ray_origin, vertex_out.structure_id, max(0., intersection_t.x));
    let world_hit_pos = hit.hit_pos + vertex_out.chunk_position;
    let clip = mul(constants.view_proj, float4(world_hit_pos, 1.));
    let depth = clip.z / clip.w;

    var frag_out : FragmentOutput;
    frag_out.depth = depth;
    frag_out.color = hit.color;

    return frag_out;
}

bool3 min_mask(float3 val)
{
    let min_val = min(val.x, min(val.y, val.z));
    return val == min_val;
}

[[vk::binding(0, 1)]] 
[[format("rgba32f")]]
WTexture2D<float4> output_texture;

[shader("compute")]
[numthreads(16, 16, 1)]
void ray_march_compute(uint32_t3 location : SV_DispatchThreadID)
{
    if(location.x >= constants.width || location.y >= constants.height)
    {
        return;
    }

    uint32_t2 pixel_loc = uint32_t2(location.x, location.y);

    let ndc_loc_x = (float)location.x / (float)constants.width * 2. - 1.;
    let ndc_loc_y = 1. - (float)location.y / (float)constants.height * 2.;

    var world_loc = mul(constants.view_proj, float4(ndc_loc_x, ndc_loc_y, 1., 1.));
    world_loc /= world_loc.w;
    let ray_direction = normalize(world_loc.xyz - constants.cam_pos);

    var inter = box_intersect(constants.cam_pos, ray_direction, float3(0.), float3(constants.chunk_width, constants.chunk_alt, constants.chunk_height));
    if(inter.y <= inter.x || inter.y <= 0.)
    {
        output_texture.Store(pixel_loc, float4(0.));
        return;
    }
    inter.x = max(0., inter.x);


    let start_position = inter.x * ray_direction + constants.cam_pos;
    // FVT
    int32_t3 current_voxel = clamp(
        int32_t3(floor(start_position)), 
        int32_t3(0),
        int32_t3(constants.chunk_width - 1, constants.chunk_alt - 1, constants.chunk_height - 1)
    );
    int32_t3 offset = int32_t3(sign(ray_direction));
    float3 delta = abs(1. / ray_direction);
    float3 position = start_position;
    float3 t = select(ray_direction > 0., current_voxel + int32_t3(1) - start_position, start_position - current_voxel) * delta;

    var color = float4(0.);
    var hit_pos = float3(0.);
    for(int32_t i = 0; i < 256; i++)
    {
        let min_mask = min_mask(t);    
        let t_adv = select(min_mask, t, float3(0.));
        let t_ray = t_adv.x + t_adv.y + t_adv.z;


        let chunk_index = current_voxel.y + current_voxel.z * constants.chunk_alt + current_voxel.x * constants.chunk_alt * constants.chunk_height;
        let hit = ray_march(ray_direction, position - float3(current_voxel), chunk_index, length(position - constants.cam_pos));
        position = start_position + ray_direction * t_ray;
        if(hit.hit)
        {
            color = hit.color;
            break;
        }
        t += select(min_mask, delta, float3(0.));
        current_voxel += select(min_mask, offset, int32_t3(0));

        if(any(current_voxel < 0) || any(current_voxel >= int32_t3(constants.chunk_width, constants.chunk_alt, constants.chunk_height)))
        {
            color = float4(0.);
            break;
        }
    }
    output_texture.Store(pixel_loc, float4(color));
}

float2 box_intersect(float3 origin, float3 ray_direction, float3 box_min, float3 box_max)
{
    let min_ts = (box_min - origin) / ray_direction;
    let max_ts = (box_max - origin) / ray_direction;

    let far_ts = max(min_ts, max_ts);
    let near_ts = min(min_ts, max_ts);

    let far_t = min(far_ts.x, min(far_ts.y, far_ts.z));
    let near_t = max(near_ts.x, max(near_ts.y, near_ts.z));
    return float2(near_t, far_t);
}