struct VertexOutput
{
    @builtin(position) postion: vec4<f32>,
    @location(0) @interpolate(flat) chunk_index: u32,
    @location(1) color: vec4<f32>,
    @location(2) cam_pos: vec3<f32>,
    @location(3) world_pos: vec3<f32>,
    @location(4) @interpolate(flat) structure_id: u32,
    @location(5) chunk_position: vec3<f32>
}

struct ChunkImmediate
{
    view_proj: mat4x4<f32>,
    cam_pos: vec3<f32>,
    frame_timestamp: u32,
}

var<immediate> constants: ChunkImmediate;
//var<push_constant> constants: ChunkInfo;

struct CacheChunkObject
{
    transform: mat4x4<f32>,
    color: vec4<f32>,
    id: u32,
    pointer: u32
}


struct StructurePoolElement
{
    pointers: array<u32, 64>
}

struct RequestBufferElement
{
    requests: array<atomic<u32>, 64>
}

struct ColorPoolElement
{
    colors: array<u32, 64>
}

struct LocationPoolElement
{
    structure_id: u32,
    structure_locator: u32
}

struct SortedRequestsElement
{
    node: u32,
    child: u32
}

fn unpack_color(color: u32) -> vec4<f32>
{
    return vec4<f32>(
            f32(color & 0xFF) / 255.,
            f32((color >> 8) & 0xFF) / 255.,
            f32((color >> 16) & 0xFF) / 255.,
        f32((color >> 24) & 0xFF) / 255.
    );
}

@group(0) @binding(0) var<storage, read_write> structure_pool: array<StructurePoolElement>;
@group(0) @binding(1) var<storage, read_write> color_pool: array<ColorPoolElement>;
@group(0) @binding(2) var<storage, read_write> location_pool: array<LocationPoolElement>;
@group(0) @binding(3) var<storage, read_write> request_buffer: array<RequestBufferElement>;
@group(0) @binding(4) var<storage, read_write> usage_buffer: array<atomic<u32>>;
@group(0) @binding(5) var<storage, read_write> structure_table_pointer: array<u32>;
@group(0) @binding(6) var<storage, read_write> structure_table_request_buffer: array<atomic<u32>>;

struct FragmentOutput {
    @location(0) color: vec4<f32>,
    @builtin(frag_depth) depth: f32, // Equivalent to gl_FragDepth
}

@vertex
fn chunk(@builtin(vertex_index) index: u32, @location(0) position: vec3<f32>, @location(1) id: u32) -> VertexOutput
{
    let cube_vertices = array<vec3<f32>, 8>(
        vec3<f32>(0., 0., 0.),
        vec3<f32>(0., 0., 1.),
        vec3<f32>(1., 0., 1.),
        vec3<f32>(1., 0., 0.),

        vec3<f32>(0., 1., 0.),
        vec3<f32>(0., 1., 1.),
        vec3<f32>(1., 1., 1.),
        vec3<f32>(1., 1., 0.),
    );

    let cube_faces = array<u32, 24>(
        // 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 = array<u32, 6>(
        0, 1, 2, 1, 3, 2
    );


    let vertex = cube_vertices[cube_faces[quad_index * 4 + triangle_map[triangle_index]]];
    let output_vertex = constants.view_proj * vec4<f32>(vertex + position, 1.0f);

    var output: VertexOutput;
    output.postion = output_vertex;
    output.color = vec4(1.);
    output.chunk_index = 0;
    output.cam_pos = constants.cam_pos;
    output.world_pos = vertex + position;
    output.structure_id = id;
    output.chunk_position = position;

    //let output = vec4<f32>(vertex, 1.0f);
    return output;
}


struct StructureElement
{
    children: array<u32, 64> 
}

struct ColorElement 
{
    children: array<vec4<f32>, 64> 
}

struct LocationElement 
{
    children: array<vec4<f32>, 64> 
}

struct RequestElement
{
    children: array<atomic<u32>, 64> 
}

fn box_inter(pos: vec3<f32>, ray_dir: vec3<f32>, box_min: vec3<f32>, box_max: vec3<f32>) -> vec2<f32>
{
    let box_min_t = (box_min - pos) / ray_dir;
    let box_max_t = (box_max - pos) / ray_dir;

    let near_ts = min(box_min_t, box_max_t);
    let far_ts = max(box_min_t, box_max_t);

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

    return vec2(near_t, far_t);
}

fn sdf(voxel: vec3<i32>) -> bool
{
    let len = length(vec3<f32>(voxel) - vec3(128)) / 128.;
    return len <= 1.;
}

fn min_vec(x: vec3<f32>) -> f32
{
    return min(x.x, min(x.y, x.z));
}

fn min_mask(x: vec3<f32>) -> vec3<bool>
{
    let min = min(x.x, min(x.y, x.z));

    return vec3<bool>(min == x.x, min == x.y, min == x.z);
}

fn node_subdivided(node: u32) -> bool
{
    return ((node >> 31) & 1) != 0;
}

fn node_pointer_valid(node: u32) -> bool
{
    return ((node >> 30) & 1) != 0;
}

fn node_pointer(node: u32) -> u32
{
    return node & 0x3FFFFFFF;
}

fn voxel_from_wall(position: vec3<f32>, ray_dir: vec3<f32>) -> vec3<i32>
{
    let integers = round(position);
    let wall_mask = min_mask(abs(position - vec3<f32>(integers)));
    let offsets = select(vec3<f32>(-0.5), vec3<f32>(0.5), ray_dir > vec3(0.));
    return vec3<i32>(floor(position + select(vec3<f32>(0.), offsets, wall_mask)));
}

struct HitResult
{
    color: vec4<f32>,
    hit_pos: vec3<f32>
}

fn new_traverse(ray_dir: vec3<f32>, ray_origin: vec3<f32>, root_id: u32, dist_offset: f32) -> HitResult
{
    let max_depth = 5;
    let dist_offset_voxel = dist_offset * f32(1 << u32(max_depth * 2));
    let fovy_deg = 100.;
    let fovy_rad = (fovy_deg / 180.) * 3.14159;
    let cone_factor = tan(fovy_rad) * 2.;

    let st_pointer = structure_table_pointer[root_id];

    
    if (!node_subdivided(st_pointer))
    {
        var result: HitResult;
        result.color = vec4(0., 1., 0., 1.);
        result.hit_pos = ray_origin;
        return result;
    }
    if(!node_pointer_valid(st_pointer))
    {
        // Node is subdivided, but not valid
        // Send request on structure table
        atomicAdd(&structure_table_request_buffer[root_id], 1); 

        var result: HitResult;
        result.color = vec4(0., 1., 0., 1.);
        result.hit_pos = ray_origin;
        return result;
    }
    //var current_node = node_pointer(st_pointer);

    // Record usage
    var dfs_stack = array<u32, 6>(node_pointer(st_pointer), 0, 0, 0, 0, 0);
    var current_depth = 0;

    usage_buffer[dfs_stack[current_depth]] = constants.frame_timestamp;

    // Start location
    //let voxel_dir = select(vec3(-1), vec3(1), ray_dir >= vec3(0.));
    var node_shift = (max_depth - current_depth) * 2;

    var child_size = 1 << u32(node_shift - 2);
    var node_size = 1 << u32(node_shift);

    var pos_origin = clamp(ray_origin * f32(1 << u32(max_depth * 2)), vec3(0.), vec3(f32(node_size) - 1.));
    var voxel = vec3<i32>(pos_origin);
    var far_t = 0.;
    var inv_ray_dir = 1. / ray_dir;
    var ray_positive = ray_dir > vec3(0.);
    var step_dir = select(vec3(-1), vec3(1), ray_positive);

    for(var iter = 0; iter < 400; iter ++)
    {
        // Compute child position
        node_shift = (max_depth - current_depth) * 2;
        child_size = 1 << u32(node_shift - 2);

        var child_pos = (voxel >> vec3(u32(node_shift - 2))) & vec3(3);
        var pointer = structure_pool[dfs_stack[current_depth]].pointers[child_pos.x + child_pos.y * 4 + child_pos.z * 4 * 4];
        
        let min_child_size = (length(vec3<f32>(voxel) - pos_origin) + dist_offset) * cone_factor;
        while(node_subdivided(pointer) && 
            f32(child_size / 8) >= min_child_size
            )
        {
            if(!node_pointer_valid(pointer) && node_subdivided(pointer))
            {
                // Record request
                atomicAdd(&request_buffer[dfs_stack[current_depth]].requests[child_pos.x + child_pos.y * 4 + child_pos.z * 4 * 4], 1);
                break;
            }

            // Descend 
            current_depth += 1;

            node_shift = (max_depth - current_depth) * 2;
            child_size = 1 << u32(node_shift - 2);
            child_pos = (voxel >> vec3(u32(node_shift - 2))) & vec3(3);
            dfs_stack[current_depth] = node_pointer(pointer);

            pointer = structure_pool[dfs_stack[current_depth]].pointers[child_pos.x + child_pos.y * 4 + child_pos.z * 4 * 4];

            // Record usage
            usage_buffer[dfs_stack[current_depth]] = constants.frame_timestamp;
        }


        // Check color
        let color = color_pool[dfs_stack[current_depth]].colors[child_pos.x + child_pos.y * 4 + child_pos.z * 4 * 4];
        if(((color >> 24) & 0xFF) != 0)
        {
            var result: HitResult;
            result.color = unpack_color(color);
            result.hit_pos = (far_t / f32(1 << u32(max_depth * 2))) * ray_dir + ray_origin;
            return result;
        }

        // Advance
        child_pos = voxel & vec3(i32(0xFFFFFFFF << u32(node_shift - 2)));
        let far_wall = child_pos + select(vec3(0), vec3(child_size), ray_positive);
        let far_wall_inter = (vec3<f32>(far_wall) - pos_origin) * inv_ray_dir;
        far_t = min(min(far_wall_inter.x, far_wall_inter.y), far_wall_inter.z);

        // Perform dda step on the children scale
        //let next_child = select(child_pos, child_pos + select(vec3(-1), vec3(1), ray_dir > vec3(0.)) * vec3(child_size), vec3(far_t) == far_wall_inter);
        let next_child = select(child_pos, child_pos + step_dir * vec3(child_size), vec3(far_t) == far_wall_inter);
        
        let previous_voxel = voxel;
        voxel = clamp(vec3<i32>(pos_origin + far_t * ray_dir), next_child, next_child + vec3(child_size) - vec3(1));

        if any(voxel < vec3(0)) || any(voxel >= vec3(1 << u32((max_depth * 2))))
        {
            discard;
        }
        
        // We touched a voxel as if we explored blocks sized by the child size of the current node.
        // But we might have exited the current node.
        
        // If this is the case we have to walk back up the tree
        // And then back down to the next node over

        // As such we find the lowest ancestor that can contain both the privous voxel (in node) and the new voxel (out of node)
        let bit_diffs = voxel ^ previous_voxel;
        let bit_diffs_lowest = bit_diffs.x | bit_diffs.y | bit_diffs.z;

        let common_depth = ((countLeadingZeros(bit_diffs_lowest) - i32(32 - max_depth * 2)) / 2);

        current_depth = common_depth;
        //current_node = dfs_stack[current_depth];
    }

    // Iter max color
    var result: HitResult;
    result.color = vec4(1., 0., 1., 1.);
    result.hit_pos = (far_t / f32(1 << u32(max_depth * 2))) * ray_dir + ray_origin;
    return result;
}

fn traverse(ray_dir: vec3<f32>, ray_origin: vec3<f32>, root_id: u32, dist_offset: f32) -> vec4<f32>
{
    let st_pointer = structure_table_pointer[root_id];

    if (!node_subdivided(st_pointer))
    {
        return vec4(0., 1., 0., 1.);
    }
    if(!node_pointer_valid(st_pointer))
    {
        atomicAdd(&structure_table_request_buffer[root_id], 1); 
        return vec4(0., 1., 0., 1.);
    }

    let fovy_deg = 100.;
    let fovy = 3.14159 * (fovy_deg / 180.);
    let definition = 1920.;
    let cone_fovy = fovy / definition; 

    let factor = 1.;
    let cone_size_factor = 2. * tan(cone_fovy) * factor;

    // Current depth of the node we are exploring
    var current_depth = 0;

    // Index of the current node's data
    var current_node = u32(st_pointer & 0x3FFFFFFF);
    usage_buffer[current_node] = constants.frame_timestamp;
    var dfs_stack = array<u32, 6>(current_node, 0, 0, 0, 0, 0);



    // Lut of the node_size per depth
    var node_size_lut = array<i32, 6>(
            4 * 4 * 4 * 4 * 4,
            4 * 4 * 4 * 4,
            4 * 4 * 4,
            4 * 4,
            4,
            1,
            );

    let local_dist_offset = dist_offset * f32(node_size_lut[0]);


    // Current node size
    var node_size = node_size_lut[0]; // 128

    // Size of a child of this node
    var child_size = node_size / 4;

    // Simple FVT
    let t_off = abs(1. / ray_dir);

    // Start location
    let voxel_dir = select(vec3(-1), vec3(1), ray_dir >= vec3(0.));
    var pos_origin = clamp(ray_origin * f32(node_size), vec3(0.), vec3(f32(node_size) - 1.));
    var voxel = vec3<i32>(pos_origin);
    var last_voxel = voxel;

    let wall_offset = select(vec3(0), vec3(1), ray_dir > vec3(0.));

    let max_depth = u32(5);
    var adaptive_depth = i32(max_depth);
    var far_t = 0.;

    let ray_dir_inv = 1. / ray_dir;
    let fma_offset = - pos_origin * ray_dir_inv;

    //let depth_limit = 3;
    for(var iter = 0; iter < 400; iter ++)
    {

        // Our ray is currently touching a voxel.
        // Descend to the lowest node that contains this voxel

        // Position of the child we are in
        var child_pos = (vec3<u32>(voxel) >> vec3<u32>((max_depth - u32(current_depth + 1)) * 2)) & vec3<u32>(3); // Hardcode for 4-tree
        var child_index = child_pos.x + child_pos.y * 4 + child_pos.z * 4 * 4;

        // Current node has been used. report
        usage_buffer[current_node] = constants.frame_timestamp;
        
        while(
            node_subdivided(structure_pool[current_node].pointers[child_index]) &&
            (local_dist_offset + far_t) * cone_size_factor < f32(node_size_lut[current_depth])
            )
        {
            if(!node_pointer_valid(structure_pool[current_node].pointers[child_index]))
            {
                atomicAdd(&request_buffer[current_node].requests[child_index], 1);
                break;
            }
            // Child node is subdivided, we go in, save position in stack
            current_node = node_pointer(structure_pool[current_node].pointers[child_index]);
            usage_buffer[current_node] = constants.frame_timestamp;
            current_depth += 1;
            dfs_stack[current_depth] = current_node;
            node_size = node_size_lut[current_depth];

            child_pos = (vec3<u32>(voxel) >> vec3<u32>((max_depth - u32(current_depth + 1)) * 2)) & vec3<u32>(3); // Hardcode for 4-tree
            child_index = child_pos.x + child_pos.y * 4 + child_pos.z * 4 * 4;
            child_size = node_size / 4;
        }
        usage_buffer[current_node] = constants.frame_timestamp;




        // At this point current_depth is the depth of the node that contains the voxel
        // child_pos and child_index relate to the specific child of the node that contains this voxel
        // It is guaranteed that the child is leave

        // Check current leave's color
        let color = unpack_color(color_pool[current_node].colors[child_index]);
        if(color.w != 0.) // Not transparent
        {
            /*
            let k = child_pos.x + child_pos.y + child_pos.z;
            let w = voxel.x + voxel.y + voxel.z;
            let x = select(0.5, 1., k % 2 == 0) * select(0.8, 1., w % 2 == 0);

            var div = 1;
            var overlay = 1.;
            for(var i = 1; i <= 5; i++)
            {
                let x = (voxel.x / div + voxel.y / div + voxel.z / div) % 2 == 0;
                overlay -= select(0., 1. / (f32(i) * 2.5), x);
                div *= 4;
            }
            */

            return color;
            //return overlay * color;
        }

        // Voxel and whole child containing it is empty
        
        // Perform a step through the children of the node
        let child_position = (voxel / child_size) * child_size;
        let far_corner = child_position + wall_offset * child_size;
        //let far_ts = (vec3<f32>(far_corner) - pos_origin) / ray_dir; // TODO: Turn into fma
        let far_ts = fma(vec3<f32>(far_corner), ray_dir_inv, fma_offset);
        far_t = min(min(far_ts.x, far_ts.y), far_ts.z);

        let next_child_min = select(child_position, child_position + voxel_dir * child_size, vec3(far_t) == far_ts);
        let next_child_max = next_child_min + vec3(child_size) - vec3(1);
        
        // The ray (far_t) is now touching the new child to explore
        // Find out which actual voxel we are touching
        let previous_voxel = voxel;
        let float_voxel = clamp(vec3<i32>(pos_origin + far_t * ray_dir), next_child_min, next_child_max);
        /*
        voxel = vec3<i32>(
            floor(
            select(
                float_voxel - vec3(0.5),
                float_voxel + vec3(0.5),
                ray_dir > vec3(0.)
            ))
        );
        */
        //voxel = vec3<i32>(round(float_voxel));
        //voxel = voxel_from_wall(float_voxel, ray_dir);
        //voxel = voxel_from_wall(float_voxel, ray_dir);
        voxel = float_voxel;
        if(any(voxel < vec3(0)) || any(voxel >= vec3(node_size_lut[0])))
        {
            //return vec4(f32(iter) / 100.);
            discard;
        }

        // We touched a voxel as if we explored blocks sized by the child size of the current node.
        // But we might have exited the current node.
        
        // If this is the case we have to walk back up the tree
        // And then back down to the next node over

        // As such we find the lowest ancestor that can contain both the privous voxel (in node) and the new voxel (out of node)
        let bit_diffs = voxel ^ previous_voxel;
        let bit_diffs_lowest = bit_diffs.x | bit_diffs.y | bit_diffs.z;

        let flb = ((countLeadingZeros(bit_diffs_lowest) - i32(32 - max_depth * 2)) / 2);
        let common_depth = flb;

        current_depth = common_depth;
        node_size = node_size_lut[current_depth];
        child_size = node_size / 4;
        current_node = dfs_stack[current_depth];
        
        // Figure out current voxel position
        //voxel = vec3<i32>(ray_origin + ray_dir * t); 

    }
    return vec4<f32>(1., 0., 1., 1.);

}

@early_depth_test(less_equal)
@fragment
fn fragment(in: VertexOutput) -> FragmentOutput
{
    //frag_out.color = vec4<f32>(2 * 0.01 / (100. + 0.01 - depth * (100. - 0.01)));
    let ray_dir = normalize(in.world_pos - in.cam_pos);
    let interp = box_inter(in.cam_pos - in.chunk_position, ray_dir, vec3(0.), vec3(1));
    let ray_origin = (in.cam_pos - in.chunk_position) + ray_dir * (max(0., interp.x));


    let result = new_traverse(ray_dir, ray_origin, in.structure_id, length(in.cam_pos - ray_origin));
    let clip_pos = constants.view_proj * vec4(result.hit_pos + in.chunk_position, 1.);
    let depth = clip_pos.z / clip_pos.w;
    var frag_out: FragmentOutput;
    frag_out.color = result.color;
    frag_out.depth = depth;
    return frag_out;

    //return vec4<f32>(ray_origin, 1.);
    //return frag_out;
    //return vec4(interp.y / 10.);
}

/*
@fragment
fn fragment(in: VertexOutput) -> @location(0) vec4<f32>
{
    let st = structure_table_pointer[0];
    let subdivided = ((st >> 31) & 1) != 0;
    let pointer_valid = ((st >> 30) & 1) != 0;
    // Request stuff
    atomicAdd(&structure_table_request_buffer[0], 1); 
    if(subdivided && !pointer_valid)
    {
        return vec4(0., 1., 0., 1.);
    }
    return vec4(1., 0., 0., 1.);
}
*/