# based on: 
# - https://www.youtube.com/watch?v=5CKvGYqagyI
# - https://pastebin.com/wXCqk1nQ

class_name PixelCountComputeShader extends Node

@export var paint_viewport_path: NodePath = NodePath()

var red_score := 0
var blue_score := 0
var pink_score := 0

var thread: Thread
var semaphore: Semaphore
var mutex: Mutex
var exit := false

var paint_texture: ViewportTexture
var paint_viewport: Viewport

var rd: RenderingDevice
var shader
var pipeline

var pixel_count_result_buffer
var pixel_count_uniform

var v_tex
var samp
var tex_uniform

var uniform_set 

func _ready():
	paint_viewport = get_node(paint_viewport_path)
	paint_texture = paint_viewport.get_texture()
	
	mutex = Mutex.new()
	semaphore = Semaphore.new()
	thread = Thread.new()
	thread.start(_thread_calculate_score)

	# We will be using our own RenderingDevice to handle the compute commands
	rd = RenderingServer.create_local_rendering_device()
	
	# Create shader and pipeline
	var shader_file = load("res://materials/shader/PixelCountCompute.glsl")
	var shader_spirv = shader_file.get_spirv()
	shader = rd.shader_create_from_spirv(shader_spirv)
	pipeline = rd.compute_pipeline_create(shader)
	
	var pb = PackedInt32Array([0,0,0])
	var pbb = pb.to_byte_array()
	
	pixel_count_result_buffer = rd.storage_buffer_create(pbb.size(), pbb)
	
	pixel_count_uniform = RDUniform.new()
	pixel_count_uniform.uniform_type = RenderingDevice.UNIFORM_TYPE_STORAGE_BUFFER
	pixel_count_uniform.binding = 0
	pixel_count_uniform.add_id(pixel_count_result_buffer)
	
	var img = paint_texture.get_image()
	var ta = img.get_data()
	
	var fmt = RDTextureFormat.new()
	fmt.width = paint_texture.get_width()
	fmt.height = paint_texture.get_height()
	fmt.usage_bits = RenderingDevice.TEXTURE_USAGE_CAN_UPDATE_BIT | RenderingDevice.TEXTURE_USAGE_SAMPLING_BIT
	fmt.format = RenderingDevice.DATA_FORMAT_R8G8B8A8_SRGB
	
	v_tex = rd.texture_create(fmt, RDTextureView.new(), [ta])
	var samp_state = RDSamplerState.new()
	samp_state.unnormalized_uvw = true
	samp = rd.sampler_create(samp_state)
	
	tex_uniform = RDUniform.new()
	tex_uniform.uniform_type = RenderingDevice.UNIFORM_TYPE_SAMPLER_WITH_TEXTURE
	tex_uniform.binding = 1
	tex_uniform.add_id(samp)
	tex_uniform.add_id(v_tex)
	
	uniform_set = rd.uniform_set_create([pixel_count_uniform, tex_uniform], shader, 0)

func recalculate_score():
	await RenderingServer.frame_post_draw
	semaphore.post()
	
func get_blue_score() -> int:
	var result := blue_score
	
	return result
	
func get_red_score() -> int:
	var result := red_score
	
	return result

func _thread_calculate_score():
	while true:
		semaphore.wait()
		
		mutex.lock()
		var should_exit := exit
		mutex.unlock()
		
		if should_exit:
			break
		
		mutex.lock()
		var image := paint_texture.get_image()
		#var image := paint_viewport.get_texture()
		mutex.unlock()
		
		var size = image.get_size()
		var ta = image.get_data()
		
		rd.texture_update(v_tex, 0, ta)
	
		# Start compute list to start recording our compute commands
		var compute_list = rd.compute_list_begin()
		# Bind the pipeline, this tells the GPU what shader to use
		rd.compute_list_bind_compute_pipeline(compute_list, pipeline)
		# Binds the uniform set with the data we want to give our shader
		rd.compute_list_bind_uniform_set(compute_list, uniform_set, 0)
		# Dispatch 1x1x1 (XxYxZ) work groups
		rd.compute_list_dispatch(compute_list, 1024/4, 1024/4, 1)
		#rd.compute_list_add_barrier(compute_list)
		# Tell the GPU we are done with this compute task
		rd.compute_list_end()

		# Force the GPU to start our commands
		rd.submit()
		# Force the CPU to wait for the GPU to finish with the recorded commands
		rd.sync()

		mutex.lock()
		# Now we can grab our data from the storage buffer
		var byte_data = rd.buffer_get_data(pixel_count_result_buffer)
		var output := byte_data.to_int32_array()
		
		red_score = floor(output[0]/100)
		blue_score = floor(output[1]/100)
		pink_score = floor(output[2]/100)
		
		var pb = PackedInt32Array([0,0,0])
		var pbb = pb.to_byte_array()
		rd.buffer_update(pixel_count_result_buffer, 0, pbb.size(), pbb)
		mutex.unlock()

func _exit_tree():
	mutex.lock()
	exit = true
	mutex.unlock()
	
	semaphore.post()
	thread.wait_to_finish()