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

#define GLSL
#include "ray_interop.h"
#undef GLSL

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

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

layout(set = 0, binding = 2) uniform UBO { UniformBuffer ubo; } ubo;
layout(set = 0, binding = 6) uniform sampler2D textures[MAX_TEXTURES];
layout(set = 0, binding = 7) uniform samplerCube skybox;

layout(set = 0, binding = 10, rgba8) uniform writeonly image2D out_base_color_a;
layout(set = 0, binding = 11, rgba32f) uniform writeonly image2D out_position_t;
layout(set = 0, binding = 12, rgba16f) uniform writeonly image2D out_normals_gs;
layout(set = 0, binding = 13, rgba8) uniform writeonly image2D out_material_rmxx;
layout(set = 0, binding = 14, rgba16f) uniform writeonly image2D out_emissive;
layout(set = 0, binding = 15, rgba32f) uniform writeonly image2D out_geometry_prev_position;
layout(set = 0, binding = 16, rgba16f) uniform writeonly image2D out_legacy_blend;

layout(set = 0, binding = 30, std430) readonly buffer ModelHeaders { ModelHeader a[]; } model_headers;
layout(set = 0, binding = 31, std430) readonly buffer Kusochki { Kusok a[]; } kusochki;
layout(set = 0, binding = 32, std430) readonly buffer Indices { uint16_t a[]; } indices;
layout(set = 0, binding = 33, std430) readonly buffer Vertices { Vertex a[]; } vertices;

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

#include "trace_simple_blending.glsl"
#include "skybox.glsl"

struct Ray {
	vec3 origin, direction;
	float dist;
};

vec3 clipToWorldSpace(vec3 clip) {
	const vec4 eye_space = ubo.ubo.inv_proj * vec4(clip, 1.);
	return (ubo.ubo.inv_view * vec4(eye_space.xyz / eye_space.w, 1.)).xyz;
}

Ray getPrimaryRay(in vec2 uv) {
	uv = uv * 2. - 1.;
	const vec3 world_near = clipToWorldSpace(vec3(uv, 0.));
	const vec3 world_far = clipToWorldSpace(vec3(uv, 1.));

	Ray ret;
	ret.origin = world_near;
	ret.direction = world_far - world_near;
	ret.dist = length(ret.direction);
	ret.direction /= ret.dist;
	return ret;
}

void main() {
	const ivec2 pix = ivec2(gl_GlobalInvocationID);
	const ivec2 res = ubo.ubo.res;
	if (any(greaterThanEqual(pix, res))) {
		return;
	}

	const vec2 uv = (gl_GlobalInvocationID.xy + .5) / res;
	const Ray ray = getPrimaryRay(uv);

	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
		;
	rayQueryInitializeEXT(rq, tlas, flags, GEOMETRY_BIT_OPAQUE | GEOMETRY_BIT_ALPHA_TEST | GEOMETRY_BIT_REFRACTIVE, ray.origin, 0., ray.direction, ray.dist);
	while (rayQueryProceedEXT(rq)) {
		// FIXME this is a no-op. It doesn't do what I though it did. Should check for SBT index for alpha-test material instead.
		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).material.tex_base_color;
		// Should never happen: skybox is opaque if (tex_base_color == TEX_BASE_SKYBOX)
		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);
		}
	}

	float L = ray.dist;
	//uint debug_geometry_index = 0;
	if (rayQueryGetIntersectionTypeEXT(rq, true) == gl_RayQueryCommittedIntersectionTriangleEXT) {
		//debug_geometry_index = rayQueryGetIntersectionGeometryIndexEXT(rq, true);
		//debug_geometry_index = rayQueryGetIntersectionPrimitiveIndexEXT(rq, true);
		primaryRayHit(rq, payload);
		L = rayQueryGetIntersectionTEXT(rq, true);
	} else {
		// Draw skybox when nothing is hit
		payload.emissive.rgb = sampleSkybox(ray.direction);
	}

	const vec4 blend = traceLegacyBlending(ray.origin, ray.direction, L);
	imageStore(out_legacy_blend, pix, blend);

	imageStore(out_position_t, pix, payload.hit_t);
	imageStore(out_base_color_a, pix, LINEARtoSRGB(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);
}