enable wgpu_mesh_shader;

struct Camera {
	// from camera to screen
	proj: mat4x4<f32>,
	// from screen to camera
	proj_inv: mat4x4<f32>,
	// from world to camera
	view: mat4x4<f32>,
	// from camera to world
	view_inv: mat4x4<f32>,
};

//group 0 is the camera
@group(0)
@binding(0)
var<uniform> camera: Camera;

struct SkyOutput {
	@builtin(position) position: vec4<f32>,
	@location(0) sampledir: vec3<f32>,
};

@vertex
fn vs_sky(@builtin(vertex_index) vertex_index: u32) -> SkyOutput {
	// hacky way to draw a large triangle
	let tmp1 = i32(vertex_index) / 2;
	let tmp2 = i32(vertex_index) & 1;
	let pos = vec4<f32>(
		f32(tmp1) * 4.0 - 1.0,
		f32(tmp2) * 4.0 - 1.0,
		1.0,
		1.0
	);

	let inv_model_view = mat3x3<f32>(camera.view_inv[0].xyz, camera.view_inv[1].xyz, camera.view_inv[2].xyz);
	let unprojected = camera.proj_inv * pos;

	var result: SkyOutput;
	result.sampledir = inv_model_view * unprojected.xyz;
	result.position = pos;
	return result;
}

struct ModelInstance{
	transform:mat4x4<f32>,
	normal_transform:mat3x3<f32>,
	color:vec4<f32>,
}
//my fancy idea is to create a megatexture for each model that includes all the textures each intance will need
//the texture transform then maps the texture coordinates to the location of the specific texture
//group 1 is the model
const MAX_MODEL_INSTANCES=512;
@group(2)
@binding(0)
var<uniform> model_instances: array<ModelInstance, MAX_MODEL_INSTANCES>;
@group(2)
@binding(1)
var model_texture: texture_2d<f32>;
@group(2)
@binding(2)
var model_sampler: sampler;

struct EntityOutputTexture {
	@builtin(position) position: vec4<f32>,
	@location(1) texture: vec2<f32>,
	@location(2) normal: vec3<f32>,
	@location(3) view: vec3<f32>,
	@location(4) color: vec4<f32>,
	@location(5) @interpolate(flat) model_color: vec4<f32>,
};
@vertex
fn vs_entity_texture(
	@builtin(instance_index) instance: u32,
	@location(0) pos: vec3<f32>,
	@location(1) texture: vec2<f32>,
	@location(2) normal: vec3<f32>,
	@location(3) color: vec4<f32>,
) -> EntityOutputTexture {
	var position: vec4<f32> = model_instances[instance].transform * vec4<f32>(pos, 1.0);
	var result: EntityOutputTexture;
	result.normal = model_instances[instance].normal_transform * normal;
	result.texture = texture;
	result.color = color;
	result.model_color = model_instances[instance].color;
	result.view = position.xyz - camera.view_inv[3].xyz;//col(3)
	result.position = camera.proj * camera.view * position;
	return result;
}

@group(1)
@binding(0)
var<uniform> model_instance: ModelInstance;

@group(1)
@binding(1)
var<uniform> ve_verts: array<vec4<f32>, 2>;

struct DebugEntityOutput {
	@builtin(position) position: vec4<f32>,
};

@vertex
fn vs_debug(
	@location(0) pos: vec3<f32>,
) -> DebugEntityOutput {
	var position: vec4<f32> = model_instance.transform * vec4<f32>(pos, 1.0);
	var result: DebugEntityOutput;
	result.position = camera.proj * camera.view * position;
	return result;
}

struct TaskPayload {
    four_bytes: u32,
}

var<task_payload> taskPayload: TaskPayload;

@task
@payload(taskPayload)
@workgroup_size(1)
fn ts_main() -> @builtin(mesh_task_size) vec3<u32> {
	taskPayload.four_bytes = 0;
	return vec3(1, 1, 1);
}

struct VertexOutput {
	@builtin(position) position: vec4<f32>,
	@location(0) color: vec4<f32>,
}
struct PrimitiveOutput {
	@builtin(triangle_indices) indices: vec3<u32>,
}

struct CircleOutput {
	@builtin(vertices) vertices: array<VertexOutput, 24>,
	@builtin(primitives) primitives: array<PrimitiveOutput, 22>,
	@builtin(vertex_count) vertex_count: u32,
	@builtin(primitive_count) primitive_count: u32,
}

var<workgroup> mesh_output: CircleOutput;

const tau: f32 = 3.141592653589793 * 2.0;

fn modulo(value:u32,modulus:u32)->u32{
	return value-value/modulus*modulus;
}

@mesh(mesh_output)
@payload(taskPayload)
@workgroup_size(1)
fn ms_debug_vert(){
	// circle with 24 vertices.
	const LAYERS:u32 = 3;
	const N:u32 = 3*(1<<LAYERS);
	mesh_output.vertex_count = N;
	mesh_output.primitive_count = N-2;

	var vertex_world_position: vec4<f32> = model_instance.transform * ve_verts[0];
	var vertex_screen_position: vec4<f32> = camera.proj * camera.view * vertex_world_position;

	for (var i:u32 = 0; i<N/4; i++){
		// draw a 1 unit redius circle
		var theta: f32 = f32(i) * tau / f32(N);
		var cos_sin: vec2<f32> = vec2(cos(theta), sin(theta));
		var offset: vec2<f32> = 0.5 * cos_sin;
		mesh_output.vertices[i].position = vertex_screen_position + vec4<f32>(offset, 0.0, 0.0);
		mesh_output.vertices[i+N/4].position = vertex_screen_position + vec4<f32>(-offset.y, offset.x, 0.0, 0.0);
		mesh_output.vertices[i+N/4*2].position = vertex_screen_position + vec4<f32>(-offset, 0.0, 0.0);
		mesh_output.vertices[i+N/4*3].position = vertex_screen_position + vec4<f32>(offset.y, -offset.x, 0.0, 0.0);
	}

	// max area triangle indices

	// the big triangle
	mesh_output.primitives[0].indices = vec3<u32>(0, N/3, N/3*2);

	// 3 layers of infill triangles to approximate circle better than 1 triangle.
	// we start on the outer layer because it's easier to construct this way
	var count:u32=N;
	var base:u32=1;
	for (var layer:u32 = 0; layer<LAYERS; layer++){
		count=count>>1;
		var step:u32=N/count;
		for (var i:u32 = 0; i<count; i++){
			mesh_output.primitives[base+i].indices = vec3<u32>(i*step, i*step+(step>>1), modulo(i*step+step,N));
		}
		base+=count;
	}
}

@mesh(mesh_output)
@payload(taskPayload)
@workgroup_size(1)
fn ms_debug_edge(){
	// draw two circles for now.
	const LAYERS:u32 = 3;
	const N:u32 = 2*(1<<LAYERS);
	mesh_output.vertex_count = 2*(3+2+4+8);
	mesh_output.primitive_count = 2*(1+2+4+8)+2;

	var v0_world_position: vec4<f32> = model_instance.transform * ve_verts[0];
	var v1_world_position: vec4<f32> = model_instance.transform * ve_verts[1];

	var v0_screen_position: vec4<f32> = camera.proj * camera.view * v0_world_position;
	var v1_screen_position: vec4<f32> = camera.proj * camera.view * v1_world_position;

	var edge_dir_world: vec4<f32> = normalize(v0_world_position - v1_world_position);
	var edge_dir_screen: vec4<f32> = camera.proj * camera.view * edge_dir_world;

	for (var i:u32 = 0; i<=N/2; i++){
		// two half circles that make a whole
		var theta: f32 = f32(i) * tau / f32(N);
		var cos_sin: vec2<f32> = vec2(cos(theta), sin(theta));

		// construct basis vectors
		var y_axis: vec2<f32> = edge_dir_screen.xy;
		var x_axis: vec2<f32> = y_axis.yx;
		x_axis.x = -x_axis.x;

		var offset: vec4<f32> = vec4<f32>(0.5 * (x_axis * cos_sin.x + y_axis * cos_sin.y), 0.0, 0.0);;

		mesh_output.vertices[i].position = v0_screen_position + offset;
		mesh_output.vertices[N/2+1+i].position = v1_screen_position - offset;
	}

	// max area triangle indices
	// number of primitives per circle half
	const P:u32 = N/2;
	// the big triangles between the circles
	mesh_output.primitives[0].indices = vec3<u32>(0, N/2+1, P);
	mesh_output.primitives[P].indices = vec3<u32>(N/2+1, 0, P + N/2+1);

	// 3 layers of infill triangles to approximate circle better than 1 triangle.
	// we start on the outer layer because it's easier to construct this way
	var count:u32=P;
	var base:u32=1;
	for (var layer:u32 = 0; layer<LAYERS; layer++){
		count=count>>1;
		var step:u32=P/count;
		for (var i:u32 = 0; i<count; i++){
			var indices = vec3<u32>(i*step, i*step+(step>>1), i*step+step);
			mesh_output.primitives[base+i].indices = indices;
			mesh_output.primitives[P+base+i].indices = indices + N/2+1;
		}
		base+=count;
	}
}

//group 2 is the skybox texture
@group(1)
@binding(0)
var cube_texture: texture_cube<f32>;
@group(1)
@binding(1)
var cube_sampler: sampler;

@fragment
fn fs_sky(vertex: SkyOutput) -> @location(0) vec4<f32> {
	return textureSample(cube_texture, cube_sampler, vertex.sampledir);
}

@fragment
fn fs_entity_texture(vertex: EntityOutputTexture) -> @location(0) vec4<f32> {
	let incident = normalize(vertex.view);
	let normal = normalize(vertex.normal);
	let d = dot(normal, incident);
	let reflected = incident - 2.0 * d * normal;

	let fragment_color = textureSample(model_texture, model_sampler, vertex.texture)*vertex.color;
	let reflected_color = textureSample(cube_texture, cube_sampler, reflected).rgb;
	return mix(vec4<f32>(vec3<f32>(0.05) + 0.2 * reflected_color,1.0),mix(vertex.model_color,vec4<f32>(fragment_color.rgb,1.0),fragment_color.a),0.5+0.5*abs(d));
}

@fragment
fn fs_debug(vertex: DebugEntityOutput) -> @location(0) vec4<f32> {
	return model_instance.color;
}
