#version 460 core
#extension GL_GOOGLE_include_directive : require
#extension GL_EXT_nonuniform_qualifier : enable
#extension GL_EXT_ray_query: require

#define RAY_QUERY
layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;

#include "ray_primary_common.glsl"
#include "ray_primary_hit.glsl"

#define RAY_PRIMARY_OUTPUTS(X) \
	X(10, base_color_a, rgba8) \
	X(11, position_t, rgba32f) \
	X(12, normals_gs, rgba16f) \
	X(13, material_rmxx, rgba8) \
	X(14, emissive, rgba16f) \
	X(15, geometry_prev_position, rgba32f) \

#define X(index, name, format) layout(set=0,binding=index,format) uniform writeonly image2D out_##name;
RAY_PRIMARY_OUTPUTS(X)
#undef X

layout(set = 0, binding = 1) uniform accelerationStructureEXT tlas;

#include "trace_additive.glsl"

void main() {
	const ivec2 pix = ivec2(gl_GlobalInvocationID);
	const ivec2 res = ivec2(imageSize(out_position_t));
	if (any(greaterThanEqual(pix, res))) {
		return;
	}
	const vec2 uv = (gl_GlobalInvocationID.xy + .5) / res * 2. - 1.;

	// FIXME start on a near plane
	const vec3 origin    = (ubo.ubo.inv_view * vec4(0, 0, 0, 1)).xyz;
	const vec4 target    = ubo.ubo.inv_proj * vec4(uv.x, uv.y, 1, 1);
	const vec3 direction = normalize((ubo.ubo.inv_view * vec4(target.xyz, 0)).xyz);

	RayPayloadPrimary payload;
	payload.hit_t = vec4(0.);
	payload.prev_pos_t = vec4(0.);
	payload.base_color_a = vec4(0.);
	payload.normals_gs = vec4(0.);
	payload.material_rmxx = vec4(0.);
	payload.emissive = vec4(0.);

	rayQueryEXT rq;
	const uint flags = 0
		| gl_RayFlagsCullFrontFacingTrianglesEXT
		//| gl_RayFlagsOpaqueEXT
		//| gl_RayFlagsTerminateOnFirstHitEXT
		//| gl_RayFlagsSkipClosestHitShaderEXT
		;
	float L = 10000.; // TODO Why 10k?
	rayQueryInitializeEXT(rq, tlas, flags, GEOMETRY_BIT_OPAQUE | GEOMETRY_BIT_ALPHA_TEST | GEOMETRY_BIT_REFRACTIVE, origin, 0., direction, L);
	while (rayQueryProceedEXT(rq)) {
		if (0 != (rayQueryGetRayFlagsEXT(rq) & gl_RayFlagsOpaqueEXT))
			continue;

		// alpha test
		// TODO check other possible ways of doing alpha test. They might be more efficient
		// (although in this particular primary ray case it's not taht important):
		// 1. Do a separate ray query for alpha masked geometry. Reason: here we might accidentally do the expensive
		//    texture sampling for geometry that's ultimately invisible (i.e. behind walls). Also, shader threads congruence.
		//    Separate pass could be more efficient as it'd be doing the same thing for every invocation.
		// 2. Same as the above, but also with a completely independent TLAS. Why: no need to mask-check geometry for opaque-vs-alpha
		const MiniGeometry geom = readCandidateMiniGeometry(rq);
		const uint tex_base_color = getKusok(geom.kusok_index).tex_base_color;
		const vec4 texture_color = texture(textures[nonuniformEXT(tex_base_color)], geom.uv);

		const float alpha_mask_threshold = .1f;
		if (texture_color.a >= alpha_mask_threshold) {
			rayQueryConfirmIntersectionEXT(rq);
		}
	}

	if (rayQueryGetIntersectionTypeEXT(rq, true) == gl_RayQueryCommittedIntersectionTriangleEXT) {
		primaryRayHit(rq, payload);
		L = rayQueryGetIntersectionTEXT(rq, true);
	}

	payload.emissive.rgb += traceAdditive(origin, direction, L);

	imageStore(out_position_t, pix, payload.hit_t);
	imageStore(out_base_color_a, pix, payload.base_color_a);
	imageStore(out_normals_gs, pix, payload.normals_gs);
	imageStore(out_material_rmxx, pix, payload.material_rmxx);
	imageStore(out_emissive, pix, payload.emissive);
	imageStore(out_geometry_prev_position, pix, payload.prev_pos_t);
}