extends Node

@export var mask_texture: Texture2D

var shader: RID
var pipeline: RID
var output_buffer: PackedByteArray

var output_uniform: RDUniform
var texture_uniform: RDUniform


func _ready() -> void:
	RenderingServer.call_on_render_thread(init_compute_shader)


func init_compute_shader() -> void:
	# Rendering device handle
	var rd := RenderingServer.get_rendering_device()
	# Load shader
	var shader_file: RDShaderFile = load("res://utilities/compute_shader_sandbox/compute_sum.glsl")
	# Compile shader
	shader = rd.shader_create_from_spirv(shader_file.get_spirv())
	# Create pipeline
	pipeline = rd.compute_pipeline_create(shader)

	# Build output buffer uniform
	output_buffer = PackedInt32Array([0]).to_byte_array()
	var storage_buffer := rd.storage_buffer_create(output_buffer.size(), output_buffer)
	output_uniform = RDUniform.new()
	output_uniform.uniform_type = RenderingDevice.UNIFORM_TYPE_STORAGE_BUFFER
	output_uniform.binding = 0
	output_uniform.add_id(storage_buffer)

	# Build texture uniform
	texture_uniform = RDUniform.new()
	texture_uniform.uniform_type = RenderingDevice.UNIFORM_TYPE_IMAGE
	texture_uniform.binding = 1

	var tex_view := RDTextureView.new()
	var tex_fmt := RDTextureFormat.new()
	tex_fmt.width = 64
	tex_fmt.height = 64
	tex_fmt.format = RenderingDevice.DATA_FORMAT_R8_UINT
	tex_fmt.usage_bits = (
		RenderingDevice.TEXTURE_USAGE_STORAGE_BIT + RenderingDevice.TEXTURE_USAGE_SAMPLING_BIT
	)
	var image := mask_texture.get_image()
	image.convert(Image.FORMAT_R8)
	var texture := rd.texture_create(tex_fmt, tex_view, [image.get_data()])
	texture_uniform.add_id(texture)


func _process(_delta: float) -> void:
	RenderingServer.call_on_render_thread(dispatch_compute)

	# Get result
	var out_data := output_buffer.to_int32_array()
	print(out_data)


func dispatch_compute() -> void:
	# Rendering device handle
	var rd := RenderingServer.get_rendering_device()

	# Prepare shader context
	var uniform_set := rd.uniform_set_create([output_uniform, texture_uniform], shader, 0)

	# Bind context
	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_dispatch(compute_list, 32, 32, 1)
	rd.compute_list_end()