struct Params {
	max_reflections: i32,
	min_strength: f32,
	sphere_count: i32,
}

struct Sphere {
	center: vec3f,
	radius: f32,
	emit_color: vec3f,
	reflect_color: vec3f,
	glossiness: f32,
}

struct Vertex {
	@location(0) eye: vec3f,
	@location(1) world: vec3f,
	@location(2) screen: vec2f,
}

struct Varying {
	@location(0) eye: vec3f,
	@location(1) world: vec3f,
	@builtin(position) screen: vec4f,
}

@group(0) @binding(0) var<uniform> params: Params;
@group(0) @binding(1) var<storage, read> spheres: array<Sphere>;

@vertex
fn on_vertex(in: Vertex) -> Varying {
	return Varying(in.eye, in.world, vec4(in.screen, 0.0, 1.0));
}

@fragment
fn on_fragment(in: Varying) -> @location(0) vec4f {
	return trace_fragment(in);
}

fn sqr(v: vec3f) -> f32 {
        return dot(v, v);
}

var<private> pos: vec3f;
var<private> ray: vec3f;

fn to_sphere(center: vec3f, radius: f32, t: ptr<function, f32>) -> bool {
	let c = sqr(pos - center) - radius * radius;
	let b = 2 * dot(pos - center, ray);
	let a = sqr(ray);

	let D = b * b - 4 * a * c;
	if (D <= 0) {
		return false;
	}
	*t = (- b - sqrt(D)) / (2 * a);
	if (*t < 0) {
		return false;
	}
	return true;
}

fn trace_fragment(in: Varying) -> vec4f {
	seed(in.screen);

	var result = vec4(0.0, 0.0, 0.0, 0.0);
	var color = vec3(1.0, 1.0, 1.0);
	pos = in.eye;
	ray = normalize(in.world - in.eye);

	for (var k = 0; k < params.max_reflections; k++) {
		var sphere = -1;
		var t = 1.0e9;
		for (var k = 0; k < params.sphere_count; k++) {
			var t1: f32;
			if (to_sphere(spheres[k].center, spheres[k].radius, &t1) && t1 < t) {
					sphere = k;
					t = t1;
			}
		}
		if (sphere == -1) {
			break;
		}
		let s = spheres[sphere];
		pos += t * ray;
		let normal = (pos - s.center) / s.radius;
		result += vec4(color * s.emit_color * -dot(normal, ray), 0.0);
		color *= s.reflect_color;
		let diffuse = normal + rand_sphere();
		let specular = reflect(ray, normal);
		ray = normalize(mix(diffuse, specular, s.glossiness));
		if (length(color) < params.min_strength) {
			break;
		}
	}

	return clamp(result, vec4(0.0), vec4(1.0));
}

fn hash(key : u32) -> u32 {
        var v = key;
        v *= 0xb384af1bu;
        v ^= v >> 15u;
        return v;
}

var<private> rand_state: u32;

fn seed(key: vec4f) {
	let x = bitcast<u32>(key.x);
	let y = bitcast<u32>(key.y);
	rand_state = hash(hash(x) ^ y);
}

fn rand_next() -> u32 {
	rand_state = hash(rand_state);
	return rand_state;
}

fn rand_float() -> f32 {
	return f32(rand_next()) / 0x1p32;
}

fn rand_sphere() -> vec3f {
	loop {
		let v = vec3f(rand_float(), rand_float(), rand_float()) - 0.5;
		let l = length(v);
		if (length(v) <= 0.5) {
			return v / l;
		}
	}
	return vec3f(0.0); // unreachable
}