@group(0) @binding(0) var output: texture_storage_2d<rgba8unorm, write>;
@group(0) @binding(2) var<storage, read_write> rays: array<Ray>;

struct Ray {
    pos: vec3<f32>,
    dist: f32,
    dir: vec3<f32>,
    hit: u32,
}

const MAX_RAY_STEPS: i32 = 64;

fn sd_sphere(p: vec3<f32>, d: f32) -> f32 {
    return length(p) - d;
}

fn sd_box(p: vec3<f32>, b: vec3<f32>) -> f32 {
    let d = abs(p) - b;
    return min(max(d.x, max(d.y, d.z)), 0.0) + length(max(d, vec3<f32>(0.0)));
}

fn get_voxel(c: vec3<i32>) -> bool {
    let p = vec3<f32>(c) + vec3<f32>(0.5);
    let d = max(-sd_sphere(p, 7.5), sd_box(p, vec3<f32>(6.0)));
    return d < 0.0;
}

fn get_mask(side_dist: vec3<f32>) -> vec3<f32> {
    let smallest = min(side_dist.x, min(side_dist.y, side_dist.z));
    if (smallest == side_dist.x) {
        return vec3<f32>(1.0, 0.0, 0.0);
    }
    else if (smallest == side_dist.y) {
        return vec3<f32>(0.0, 1.0, 0.0);
    }
    return vec3<f32>(0.0, 0.0, 1.0);
}

fn traverse(ray: Ray) -> Ray {
    var map_pos = floor(ray.pos);
    var ray_sign = sign(ray.dir);
    var delta_dist = 1.0 / ray.dir;
    var side_dist = delta_dist * ((map_pos - ray.pos) + ray_sign * 0.5 + 0.5);
    var mask = get_mask(side_dist);

    for (var i: i32 = 0; i < MAX_RAY_STEPS; i++) {
        if (get_voxel(vec3<i32>(map_pos))) {
            let mini = delta_dist * ((map_pos - ray.pos) - ray_sign * 0.5 + 0.5);
            let d = max(mini.x, max(mini.y, mini.z));
            return Ray(ray.pos, d, ray.dir, 1);
        }

        mask = get_mask(side_dist);
        map_pos += mask * ray_sign;
        side_dist += mask * ray_sign * delta_dist;
    }

    return ray;
}


@compute @workgroup_size(8, 8, 1)
fn compute(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let img_coord = global_id.xy;
    let img_dims = textureDimensions(output);

    // This discards the extra pixels in cases where the image size isn't perfectly divisible by the kernel.xy
    if (img_coord.x >= img_dims.x || img_coord.y >= img_dims.y) {
        return;
    }

    let idx = img_coord.x + img_coord.y * img_dims.x;
    rays[idx] = traverse(rays[idx]);

    // Temp
    var color = vec4<f32>(0.5, 0.6, 0.7, 1.0);
    let ray = rays[idx];
    if (ray.hit == 1) {
        let intersect = ray.pos + ray.dir * (ray.dist + 0.0001);
        let uv3d = intersect - floor(intersect);
        color = vec4<f32>(uv3d, 1.0);
    }
    textureStore(output, img_coord, color);
}