#version 450

layout (local_size_x = 32) in;

struct Vertex {
	vec3 position;
	vec3 color;
	vec2 uv;
	vec3 normal;
	vec3 velocity;
	vec3 prevPosition;
	float inverseMass;
};

struct Face {
	uint a;
	uint b;
	uint c;
};

layout (std430, set = 0, binding = 0) buffer VertexBuffer {
	Vertex vertices[];
};
layout (std430, set = 0, binding = 1) buffer FaceBuffer {
	Face faces[];
};

layout (std140, set = 0, binding = 5) uniform Sizes {
	uint vertexCount;
	uint faceCount;
};

layout (set = 1, binding = 0) uniform CameraUniform {
    mat4 view;
    mat4 projection;
	vec2 viewport;
} camera;


layout (std430, set = 2, binding = 1) buffer GrabInformation {
	float originalInverseMass;
	uint vID;
	uint distanceToFace; // bits are float format
	bool foundHit;
};

layout (push_constant, std430) uniform PushConstants {
	uint state;
	vec2 screenPosition;
	vec2 screenDelta;
};

vec3 toNDC(vec3 world){
	mat4 realProjection = camera.projection;
	realProjection[1][1] *= -1;

	mat4 VP = realProjection * camera.view;

	vec4 inhomo = VP * vec4(world, 1);

	return vec3(inhomo.xyz / inhomo.w);
}

vec3 toWorld(vec3 ndcCube){
	mat4 realProjection = camera.projection;
	realProjection[1][1] *= -1;
	mat4 VPinv = inverse(realProjection * camera.view);

	vec4 inhomo = VPinv * vec4(ndcCube, 1);

	return vec3(inhomo.xyz / inhomo.w);
}

void rayForward(out vec3 origin, out vec3 direction){
	vec2 ndc = vec2(
		(2 * screenPosition.x / camera.viewport.x) - 1,
		1 - (2 * screenPosition.y / camera.viewport.y)
	);

	vec3 a = vec3(ndc, 0);
	vec3 b = vec3(ndc, 1);

	vec3 p1 = toWorld(a);
	vec3 p2 = toWorld(b);

	origin = p1;
	direction = normalize(p2 - p1);
}

void testFace(uint fID){
	vec3 origin;
	vec3 direction;
	rayForward(origin, direction);

	Face face = faces[fID];
	vec3 a = vertices[face.a].position;
	vec3 b = vertices[face.b].position;
	vec3 c = vertices[face.c].position;

	vec3 ab = b - a;
	vec3 ac = c - a;

	vec3 n = cross(ab, ac);
	if (dot(direction, n) < 0){
		return;
	}

	float r = dot(a - origin, n) / dot(direction, n);
	if (r < 0){
		return;
	}
	vec3 q = origin + r * direction;

	vec3 qa = a - q;
	vec3 qb = b - q;
	vec3 qc = c - q;

	float area = length(cross(ab, ac)) / 2;
	float alpha = length(cross(qb, qc)) / (2 * area);
	float beta = length(cross(qc, qa)) / (2 * area);
	float gamma = length(cross(qa, qb)) / (2 * area);

	bool isOne = abs(alpha + beta + gamma - 1) < 0.001;

	if (alpha >= 0 && alpha <= 1 && beta >= 0 && beta <= 1 && gamma >= 0 && gamma <= 1 && isOne){
		uint newDistance;
		uint expectedOriginalDistance;
		do {
			expectedOriginalDistance = distanceToFace;
			newDistance = floatBitsToUint(min(r, uintBitsToFloat(expectedOriginalDistance)));
		} while (atomicCompSwap(distanceToFace, expectedOriginalDistance, newDistance) != expectedOriginalDistance);
		if (newDistance == floatBitsToUint(r)){
			foundHit = true;
			vID = face.a;
		}
	}
}

void move(){
	vec3 ndcDelta = vec3(
		2 * screenDelta.x / camera.viewport.x,
		-2 * screenDelta.y / camera.viewport.y,
		0
	);

	vec3 ndc = toNDC(vertices[vID].position);
	vec3 ndc2 = ndc + ndcDelta;

	vec3 p2 = toWorld(ndc2);

	vec3 worldDelta = p2 - vertices[vID].position;

	vertices[vID].position += worldDelta;
}

void release(){
	vertices[vID].inverseMass = originalInverseMass;
	distanceToFace = floatBitsToUint(1.0e20);
	foundHit = false;
}

void main(){
	uint id = gl_GlobalInvocationID.x;

	switch (state){
		case 0:
			if (id < faceCount){
				testFace(id);
			}
			break;
		case 1:
			if (id == 0 && foundHit){
				originalInverseMass = vertices[vID].inverseMass;
				vertices[vID].inverseMass = 0;
			}
			break;
		case 2:
			if (id == 0 && foundHit){
				move();
			}
			break;
		case 3:
			if (id == 0 && foundHit){
				release();
			}
			break;

	}
}