@tool
class_name PostProcessShader
extends CompositorEffect


@export var vignette_colour : Color = Color(0.0, 0.0, 0.0, 0.392)
@export var vignette_pupil : float = 0.8
@export var vignette_iris : float = 0.4


# Called when this resource is constructed.
func _init():
	effect_callback_type = CompositorEffect.EFFECT_CALLBACK_TYPE_POST_TRANSPARENT
	RenderingServer.call_on_render_thread(_initialize_compute)

# System notifications, we want to react on the notification that
# alerts us we are about to be destroyed.
func _notification(what):
	if what == NOTIFICATION_PREDELETE:
		if shader.is_valid():
			# Freeing our shader will also free any dependents such as the pipeline!
			rd.free_rid(shader)

var rd : RenderingDevice
var shader : RID
var pipeline : RID

## Render thread code below

func _initialize_compute():
	rd = RenderingServer.get_rendering_device()
	if !rd: return
	
	var shader_file = load("uid://bxunu0anan320")
	var shader_spirv : RDShaderSPIRV = shader_file.get_spirv()
	shader = rd.shader_create_from_spirv(shader_spirv)
	pipeline = rd.compute_pipeline_create(shader)

# Called by the rendering thread every frame.
func _render_callback(p_effect_callback_type, p_render_data):
	if rd and p_effect_callback_type == EFFECT_CALLBACK_TYPE_POST_TRANSPARENT:
		# Get our render scene buffers object, this gives us access to our render buffers.
		# Note that implementation differs per renderer hence the need for the cast.
		var render_scene_buffers: RenderSceneBuffersRD = p_render_data.get_render_scene_buffers()
		if render_scene_buffers:

			# Loop through views just in case we're doing stereo rendering. No extra cost if this is mono.
			var view_count = render_scene_buffers.get_view_count()
			for view in range(view_count):
				# Get our render size, this is the 3D render resolution!
				var raster_size = render_scene_buffers.get_internal_size()
				if raster_size.x == 0 and raster_size.y == 0:
					return

				# We can use a compute shader here.
				var x_groups = (raster_size.x - 1) / 8 + 1
				var y_groups = (raster_size.y - 1) / 8 + 1
				var z_groups = 1
				
				# Get the RID for our color image, we will be reading from and writing to it.
				var input_image = render_scene_buffers.get_color_layer(view)
				
				# Push constant.
				var push_constant: PackedFloat32Array = PackedFloat32Array()
				push_constant.push_back(raster_size.x)
				push_constant.push_back(raster_size.y)
				push_constant.push_back(0.0)
				push_constant.push_back(0.0)
				
				push_constant.push_back(vignette_colour.r)
				push_constant.push_back(vignette_colour.g)
				push_constant.push_back(vignette_colour.b)
				push_constant.push_back(vignette_colour.a)
				
				push_constant.push_back(vignette_pupil)
				push_constant.push_back(vignette_iris)
				push_constant.push_back(0.0)
				push_constant.push_back(0.0)
				
				# Create a uniform set.
				# This will be cached; the cache will be cleared if our viewport's configuration is changed.
				var uniform: RDUniform = RDUniform.new()
				uniform.uniform_type = RenderingDevice.UNIFORM_TYPE_IMAGE
				uniform.binding = 0
				uniform.add_id(input_image)
				var uniform_set = UniformSetCacheRD.get_cache(shader, 0, [ uniform ])
				
				# Run our compute shader.
				var compute_list:= rd.compute_list_begin()
				rd.compute_list_bind_compute_pipeline(compute_list, pipeline)
				rd.compute_list_bind_uniform_set(compute_list, uniform_set, 0)
				rd.compute_list_set_push_constant(compute_list, push_constant.to_byte_array(), push_constant.size() * 4)
				rd.compute_list_dispatch(compute_list, x_groups, y_groups, z_groups)
				rd.compute_list_end()