// Copyright 2021-2023, Collabora Ltd.
// Copyright 2025, NVIDIA CORPORATION.
// Author: Jakob Bornecrantz <jakob@collabora.com>
// Author: Christoph Haag <christoph.haag@collabora.com>
// SPDX-License-Identifier: BSL-1.0

#version 460
#extension GL_GOOGLE_include_directive : require

#include "srgb.inc.glsl"
#include "layer_defines.inc.glsl"


struct layer_data
{
	uint layer_type;
	uint unpremultiplied_alpha;

	// This struct is used in an array, gets padded to vec4.
	uint _padding0;
	uint _padding1;
};

struct image_info
{
	uint color_image_index;
	uint depth_image_index;

	// This struct is used in an array, gets padded to vec4.
	uint _padding0;
	uint _padding1;
};


const float PI = acos(-1);

// Should we do timewarp.
layout(constant_id = 1) const bool do_timewarp = false;
layout(constant_id = 2) const bool do_color_correction = true;

//! This is always set by the render_resource pipeline creation code to the actual limit.
layout(constant_id = 3) const int RENDER_MAX_LAYERS = 128;
layout(constant_id = 4) const int SAMPLER_ARRAY_SIZE = 16;

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

// layer 0 color, [optional: layer 0 depth], layer 1, ...
layout(set = 0, binding = 0) uniform sampler2D source[SAMPLER_ARRAY_SIZE];
layout(set = 0, binding = 2) uniform writeonly restrict image2D target;
layout(set = 0, binding = 3, std140) uniform restrict Config
{
	ivec4 view;
	ivec4 layer_count;

	vec4 pre_transform;
	vec4 post_transform[RENDER_MAX_LAYERS];

	// Per-layer data.
	layer_data layer_data[RENDER_MAX_LAYERS];

	// which image/sampler(s) correspond to each layer
	image_info image_info[RENDER_MAX_LAYERS];

	// shared between cylinder and equirect2
	mat4 mv_inverse[RENDER_MAX_LAYERS];


	// for cylinder layer
	vec4 cylinder_data[RENDER_MAX_LAYERS];


	// for equirect2 layer
	vec4 eq2_data[RENDER_MAX_LAYERS];


	// for projection layers

	// timewarp matrices
	mat4 transform[RENDER_MAX_LAYERS];


	// for quad layers

	// all quad transforms and coordinates are in view space
	vec4 quad_position[RENDER_MAX_LAYERS];
	vec4 quad_normal[RENDER_MAX_LAYERS];
	mat4 inverse_quad_transform[RENDER_MAX_LAYERS];

	// quad extent in world scale
	vec2 quad_extent[RENDER_MAX_LAYERS];
} ubo;


vec2 position_to_view_uv(ivec2 extent, uint ix, uint iy)
{
	// Turn the index into floating point.
	vec2 xy = vec2(float(ix), float(iy));

	// The inverse of the extent of a view image is the pixel size in [0 .. 1] space.
	vec2 extent_pixel_size = vec2(1.0 / float(extent.x), 1.0 / float(extent.y));

	// Per-target pixel we move the size of the pixels.
	vec2 view_uv = xy * extent_pixel_size;

	// Emulate a triangle sample position by offset half target pixel size.
	view_uv = view_uv + extent_pixel_size / 2.0;

	return view_uv;
}

vec2 transform_uv_subimage(vec2 uv, uint layer)
{
	vec2 values = uv;

	// To deal with OpenGL flip and sub image view.
	values.xy = fma(values.xy, ubo.post_transform[layer].zw, ubo.post_transform[layer].xy);

	// Ready to be used.
	return values.xy;
}

vec2 transform_uv_timewarp(vec2 uv, uint layer)
{
	vec4 values = vec4(uv, -1, 1);

	// From uv to tan angle (tangent space).
	values.xy = fma(values.xy, ubo.pre_transform.zw, ubo.pre_transform.xy);
	values.y = -values.y; // Flip to OpenXR coordinate system.

	// Timewarp.
	values = ubo.transform[layer] * values;
	values.xy = values.xy * (1.0 / max(values.w, 0.00001));

	// From [-1, 1] to [0, 1]
	values.xy = values.xy * 0.5 + 0.5;

	// To deal with OpenGL flip and sub image view.
	values.xy = fma(values.xy, ubo.post_transform[layer].zw, ubo.post_transform[layer].xy);

	// Done.
	return values.xy;
}

vec2 transform_uv(vec2 uv, uint layer)
{
	if (do_timewarp) {
		return transform_uv_timewarp(uv, layer);
	} else {
		return transform_uv_subimage(uv, layer);
	}
}

vec4 do_cylinder(vec2 view_uv, uint layer)
{
	// Get ray position in model space.
	const vec3 ray_origin = (ubo.mv_inverse[layer] * vec4(0, 0, 0, 1)).xyz;

	// [0 .. 1] to tangent lengths (at unit Z).
	const vec2 uv = fma(view_uv, ubo.pre_transform.zw, ubo.pre_transform.xy);

	// With Z at the unit plane and flip y for OpenXR coordinate system,
	// transform the ray into model space.
	const vec3 ray_dir = normalize((ubo.mv_inverse[layer] * vec4(uv.x, -uv.y, -1, 0)).xyz);

	const float radius = ubo.cylinder_data[layer].x;
	const float central_angle = ubo.cylinder_data[layer].y;
	const float aspect_ratio = ubo.cylinder_data[layer].z;

	vec3 dir_from_cyl;
	// CPU code will set +INFINITY to zero.
	if (radius == 0) {
		dir_from_cyl = ray_dir;
	} else {
		// Find if the cylinder intersects with the ray direction
		// Inspired by Inigo Quilez
		// https://iquilezles.org/articles/intersectors/

		const vec3 axis = vec3(0.f, 1.f, 0.f);

		float card = dot(axis, ray_dir);
		float caoc = dot(axis, ray_origin);
		float a = 1.f - card * card;
		float b = dot(ray_origin, ray_dir) - caoc * card;
		float c = dot(ray_origin, ray_origin) - caoc * caoc - radius * radius;
		float h = b * b - a * c;
		if(h < 0.f) {
			// no intersection
			return vec4(0.f);
		}

		h = sqrt(h);
		vec2 distances = vec2(-b - h, -b + h) / a;

		if (distances.y < 0) {
			return vec4(0.f);
		}

		dir_from_cyl = normalize(ray_origin + (ray_dir * distances.y));
	}

	const float lon = atan(dir_from_cyl.x, -dir_from_cyl.z) / (2 * PI) + 0.5; // => [0, 1]
	// float lat = -asin(dir_from_cyl.y); // => [-π/2, π/2]
	// float y = tan(lat); // => [-inf, inf]
	// simplified: -y/sqrt(1 - y^2)
	const float y = -dir_from_cyl.y / sqrt(1 - (dir_from_cyl.y * dir_from_cyl.y)); // => [-inf, inf]

	vec4 out_color = vec4(0.f);

#ifdef DEBUG
	const int lon_int = int(lon * 1000.f);
	const int y_int = int(y * 1000.f);

	if (lon < 0.001 && lon > -0.001) {
		out_color = vec4(1, 0, 0, 1);
	} else if (lon_int % 50 == 0) {
		out_color = vec4(1, 1, 1, 1);
	} else if (y_int % 50 == 0) {
		out_color = vec4(1, 1, 1, 1);
	} else {
		out_color = vec4(lon, y, 0, 1);
	}
#endif

	const float chan = central_angle / (PI * 2.f);

	// height in radii, radius only matters for determining intersection
	const float height = central_angle * aspect_ratio;

	// Normalize [0, 2π] to [0, 1]
	const float uhan = 0.5 + chan / 2.f;
	const float lhan = 0.5 - chan / 2.f;

	const float ymin = -height / 2;
	const float ymax = height / 2;

	if (y < ymax && y > ymin && lon < uhan && lon > lhan) {
		// map configured display region to whole texture
		vec2 offset = vec2(lhan, ymin);
		vec2 extent = vec2(uhan - lhan, ymax - ymin);
		vec2 sample_point = (vec2(lon, y) - offset) / extent;

		vec2 uv_sub = fma(sample_point, ubo.post_transform[layer].zw, ubo.post_transform[layer].xy);

		uint index = ubo.image_info[layer].color_image_index;
#ifdef DEBUG
		out_color += texture(source[index], uv_sub) / 2.f;
#else

		out_color = texture(source[index], uv_sub);
#endif
	} else {
		out_color += vec4(0.f);
	}

	return out_color;
}

vec4 do_equirect2(vec2 view_uv, uint layer)
{
	// Get ray position in model space.
	const vec3 ray_origin = (ubo.mv_inverse[layer] * vec4(0, 0, 0, 1)).xyz;

	// [0 .. 1] to tangent lengths (at unit Z).
	const vec2 uv = fma(view_uv, ubo.pre_transform.zw, ubo.pre_transform.xy);

	// With Z at the unit plane and flip y for OpenXR coordinate system,
	// transform the ray into model space.
	const vec3 ray_dir = normalize((ubo.mv_inverse[layer] * vec4(uv.x, -uv.y, -1, 0)).xyz);

	const float radius = ubo.eq2_data[layer].x;
	const float central_horizontal_angle = ubo.eq2_data[layer].y;
	const float upper_vertical_angle = ubo.eq2_data[layer].z;
	const float lower_vertical_angle = ubo.eq2_data[layer].w;

	vec3 dir_from_sph;
	// CPU code will set +INFINITY to zero.
	if (radius == 0) {
		dir_from_sph = ray_dir;
	} else {
		// Find if the sphere intersects with the ray using Pythagoras'
		// theroem with a triangle formed by QC, H and the radius.
		// Inspired by Inigo Quilez
		// https://iquilezles.org/articles/intersectors/

		const float B = dot(ray_origin, ray_dir);
		// QC is the point where the ray passes closest
		const vec3 QC = ray_origin - B * ray_dir;
		// If the distance is father than the radius, no hit
		float H = radius * radius - dot(QC, QC);
		if (H < 0.0) {
			// no intersection
			return vec4(0.f);
		}

		H = sqrt(H);

		vec2 distances = vec2(-B - H, -B + H);
		if (distances.y < 0) {
			return vec4(0.f);
		}

		dir_from_sph = normalize(ray_origin + (ray_dir * distances.y));
	}

	const float lon = atan(dir_from_sph.x, -dir_from_sph.z) / (2 * PI) + 0.5;
	const float lat = acos(dir_from_sph.y) / PI;

	vec4 out_color = vec4(0.f);

#ifdef DEBUG
	const int lon_int = int(lon * 1000.f);
	const int lat_int = int(lat * 1000.f);

	if (lon < 0.001 && lon > -0.001) {
		out_color = vec4(1, 0, 0, 1);
	} else if (lon_int % 50 == 0) {
		out_color = vec4(1, 1, 1, 1);
	} else if (lat_int % 50 == 0) {
		out_color = vec4(1, 1, 1, 1);
	} else {
		out_color = vec4(lon, lat, 0, 1);
	}
#endif

	const float chan = central_horizontal_angle / (PI * 2.0f);

	// Normalize [0, 2π] to [0, 1]
	const float uhan = 0.5 + chan / 2.0f;
	const float lhan = 0.5 - chan / 2.0f;

	// Normalize [-π/2, π/2] to [0, 1]
	const float uvan = upper_vertical_angle / PI + 0.5f;
	const float lvan = lower_vertical_angle / PI + 0.5f;

	if (lat < uvan && lat > lvan && lon < uhan && lon > lhan) {
		// map configured display region to whole texture
		vec2 ll_offset = vec2(lhan, lvan);
		vec2 ll_extent = vec2(uhan - lhan, uvan - lvan);
		vec2 sample_point = (vec2(lon, lat) - ll_offset) / ll_extent;

		vec2 uv_sub = fma(sample_point, ubo.post_transform[layer].zw, ubo.post_transform[layer].xy);

		uint index = ubo.image_info[layer].color_image_index;
#ifdef DEBUG
		out_color += texture(source[index], uv_sub) / 2.0;
#else

		out_color = texture(source[index], uv_sub);
#endif
	} else {
		out_color += vec4(0.f);
	}

	return out_color;
}

vec4 do_projection(vec2 view_uv, uint layer)
{
	uint source_image_index = ubo.image_info[layer].color_image_index;

	// Do any transformation needed.
	vec2 uv = transform_uv(view_uv, layer);

	// Sample the source.
	vec4 colour = vec4(texture(source[source_image_index], uv).rgba);

	return colour;
}

vec3 get_direction(vec2 uv)
{
	// Skip the DIM/STRETCH/OFFSET stuff and go directly to values
	vec4 values = vec4(uv, -1, 1);

	// From uv to tan angle (tangent space).
	values.xy = fma(values.xy, ubo.pre_transform.zw, ubo.pre_transform.xy);
	values.y = -values.y; // Flip to OpenXR coordinate system.

	// This works because values.xy are now in tangent space, that is the
	// `tan(a)` on each of the x and y axis. That means values.xyz now
	// define a point on the plane that sits at Z -1 and has a normal that
	// runs parallel to the Z-axis. So if you run normalize you get a normal
	// that points at that point.
	vec3 direction = normalize(values.xyz);

	return direction;
}

vec4 do_quad(vec2 view_uv, uint layer)
{
	uint source_image_index = ubo.image_info[layer].color_image_index;

	// center point of the plane in view space.
	vec3 quad_position = ubo.quad_position[layer].xyz;

	// normal vector of the plane.
	vec3 normal = ubo.quad_normal[layer].xyz;
	normal = normalize(normal);

	// coordinate system is the view space, therefore the camera/eye position is in the origin.
	vec3 camera = vec3(0.0, 0.0, 0.0);

	// default color white should never be visible
	vec4 colour = vec4(1.0, 1.0, 1.0, 1.0);

	//! @todo can we get better "pixel stuck" on projection layers with timewarp uv?
	// never use the timewarp uv here because it depends on the projection layer pose
	vec2 uv = view_uv;

	/*
	* To fill in the view_uv texel on the target texture, an imaginary ray is shot through texels on the target
	* texture. When this imaginary ray hits a quad layer, it means that when the respective color at the hit
	* intersection is picked for the current view_uv texel, the final image as seen through the headset will
	* show this view_uv texel at the respective location.
	*/
	vec3 direction = get_direction(uv);
	direction = normalize(direction);

	float denominator = dot(direction, normal);

	// denominator is negative when vectors point towards each other, 0 when perpendicular,
	// and positive when vectors point in a similar direction, i.e. direction vector faces quad backface, which we don't render.
	if (denominator < 0.00001) {
		// shortest distance between origin and plane defined by normal + quad_position
		float dist = dot(camera - quad_position, normal);

		// distance between origin and intersection point on the plane.
		float intersection_dist = (dot(camera, normal) + dist) / -denominator;

		// layer is behind camera as defined by direction vector
		if (intersection_dist < 0) {
			colour = vec4(0.0, 0.0, 0.0, 0.0);
			return colour;
		}

		vec3 intersection = camera + intersection_dist * direction;

		// ps for "plane space"
		vec2 intersection_ps = (ubo.inverse_quad_transform[layer] * vec4(intersection.xyz, 1.0)).xy;

		bool in_plane_bounds =
			intersection_ps.x >= - ubo.quad_extent[layer].x / 2. && //
			intersection_ps.x <= ubo.quad_extent[layer].x / 2. && //
			intersection_ps.y >= - ubo.quad_extent[layer].y / 2. && //
			intersection_ps.y <= ubo.quad_extent[layer].y / 2.;

		if (in_plane_bounds) {
			// intersection_ps is in [-quad_extent .. quad_extent]. Transform to  [0 .. quad_extent], then scale to [ 0 .. 1 ] for sampling
			vec2 plane_uv = (intersection_ps.xy + ubo.quad_extent[layer] / 2.) / ubo.quad_extent[layer];

			// sample on the desired subimage, not the entire texture
			plane_uv = fma(plane_uv, ubo.post_transform[layer].zw, ubo.post_transform[layer].xy);

			colour = texture(source[source_image_index], plane_uv);
		} else {
			// intersection on infinite plane outside of plane bounds
			colour = vec4(0.0, 0.0, 0.0, 0.0);
			return colour;
		}
	} else {
		// no intersection with front face of infinite plane or perpendicular
		colour = vec4(0.0, 0.0, 0.0, 0.0);
		return colour;
	}

	return vec4(colour);
}

vec4 do_layers(vec2 view_uv)
{
	vec4 accum = vec4(0, 0, 0, 0);

	int layer_count = ubo.layer_count.x;
	for (uint layer = 0; layer < layer_count; layer++) {
		vec4 rgba = vec4(0, 0, 0, 0);

		switch (ubo.layer_data[layer].layer_type) {
		case LAYER_COMP_TYPE_QUAD:
			rgba = do_quad(view_uv, layer);
			break;
		case LAYER_COMP_TYPE_CYLINDER:
			rgba = do_cylinder(view_uv, layer);
			break;
		case LAYER_COMP_TYPE_EQUIRECT2:
			rgba = do_equirect2(view_uv, layer);
			break;
		case LAYER_COMP_TYPE_PROJECTION:
			rgba = do_projection(view_uv, layer);
			break;
		default: break;
		}

		if (ubo.layer_data[layer].unpremultiplied_alpha != 0) {
			// Unpremultipled blend factor of src.a.
			accum.rgb = mix(accum.rgb, rgba.rgb, rgba.a);
		} else {
			// Premultiplied blend factor of 1.
			accum.rgb = (accum.rgb * (1 - rgba.a)) + rgba.rgb;
		}
		accum.a = fma((1.f - rgba.a), accum.a, rgba.a);
	}

	return accum;
}

void main()
{
	uint ix = gl_GlobalInvocationID.x;
	uint iy = gl_GlobalInvocationID.y;

	ivec2 offset = ivec2(ubo.view.xy);
	ivec2 extent = ivec2(ubo.view.zw);

	if (ix >= extent.x || iy >= extent.y) {
		return;
	}

	vec2 view_uv = position_to_view_uv(extent, ix, iy);

	vec4 colour = do_layers(view_uv);

	if (do_color_correction) {
		// Do colour correction here since there are no automatic conversion in hardware available.
		colour.rgb = from_linear_to_srgb(colour.rgb);
	}

	imageStore(target, ivec2(offset.x + ix, offset.y + iy), colour);
}