var rd: RenderingDevice
var shader_file: Resource
var shader_spirv: RDShaderSPIRV
var shader: RID

func create_device():
	rd = RenderingServer.create_local_rendering_device()
	shader_file = load("res://sdf_shader.glsl")
	shader_spirv = shader_file.get_spirv()
	shader = rd.shader_create_from_spirv(shader_spirv)

func compute(world_size: int, threshold: float):

	# Prepare our data. We use floats in the shader, so we need 32 bit.
	var total = world_size * world_size * world_size
	var input := PackedFloat32Array()
	input.resize(total)
	var input_bytes := input.to_byte_array()

	# Create a storage buffer that can hold our float values.
	var buffer := rd.storage_buffer_create(input_bytes.size(), input_bytes)

	# Create a uniform to assign the buffer to the rendering device
	var uniform_buf := RDUniform.new()
	uniform_buf.uniform_type = RenderingDevice.UNIFORM_TYPE_STORAGE_BUFFER
	uniform_buf.binding = 0 # this needs to match the "binding" in our shader file
	uniform_buf.add_id(buffer)
	
	# Prepare our data. We use floats in the shader, so we need 32 bit.
	var surface_total = world_size * world_size * world_size * 3
	var surface_input := PackedFloat32Array()
	surface_input.resize(surface_total)
	var surface_input_bytes := surface_input.to_byte_array()

	# Create a storage buffer that can hold our float values.
	var surface_buffer := rd.storage_buffer_create(surface_input_bytes.size(), surface_input_bytes)

	# Create a uniform to assign the buffer to the rendering device
	var surface_uniform_buf := RDUniform.new()
	surface_uniform_buf.uniform_type = RenderingDevice.UNIFORM_TYPE_STORAGE_BUFFER
	surface_uniform_buf.binding = 2 # this needs to match the "binding" in our shader file
	surface_uniform_buf.add_id(surface_buffer)
	
	var peer := StreamPeerBuffer.new()
	peer.put_32(world_size)
	peer.put_float(threshold)
	peer.put_32(0)
	peer.put_32(0)
	var bytes := peer.data_array

	var params_buffer := rd.uniform_buffer_create(bytes.size(), bytes)

	var uniform_params := RDUniform.new()
	uniform_params.uniform_type = RenderingDevice.UNIFORM_TYPE_UNIFORM_BUFFER
	uniform_params.binding = 1
	uniform_params.add_id(params_buffer)

	var uniform_set := rd.uniform_set_create([uniform_params, uniform_buf, surface_uniform_buf], shader, 0) # the last parameter (the 0) needs to match the "set" in our shader file
	
	# Create a compute pipeline
	var pipeline := rd.compute_pipeline_create(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)
	var dispatch_x = int(ceil(world_size / 8.0))
	var dispatch_y = int(ceil(world_size / 8.0))
	var dispatch_z = int(ceil(world_size / 1.0))  # local_size_z = 1
	rd.compute_list_dispatch(compute_list, dispatch_x, dispatch_y, dispatch_z)
	rd.compute_list_end()
	
	# Submit to GPU and wait for sync
	rd.submit()
	rd.sync()
	# Read back the data from the buffer
	var output_bytes := rd.buffer_get_data(surface_buffer)
	var output := output_bytes.to_float32_array()
	
	var output_sample_bytes := rd.buffer_get_data(buffer)
	var output_sample := output_sample_bytes.to_float32_array()
	
	rd.free_rid(buffer)
	rd.free_rid(params_buffer)
	rd.free_rid(surface_buffer)
	
	var surface_points_dict := build_surface_dict(world_size, output)
	var sample_points_dict := build_sample_dict(world_size, output_sample)
	
	return {
		"surface_points_dict": surface_points_dict,
		"sample_points_dict": sample_points_dict
	}

func build_sample_dict(world_size: int, flat_buffer: PackedFloat32Array) -> Dictionary:
	var dict := {}
	var total = world_size * world_size * world_size

	for idx in total:
		var voxel_x = idx % world_size
		var voxel_y = (idx / world_size) % world_size
		var voxel_z = idx / (world_size * world_size)

		var voxel_id = Vector3i(voxel_x, voxel_y, voxel_z)
		var distance = flat_buffer[idx]

		dict[voxel_id] = distance

	return dict
	
func build_surface_dict(world_size: int, flat_buffer: PackedFloat32Array) -> Dictionary:
	var dict := {}
	var total = world_size * world_size * world_size

	for idx in total:
		var base = idx * 3
		var x = flat_buffer[base]
		var y = flat_buffer[base + 1]
		var z = flat_buffer[base + 2]

		var voxel_x = idx % world_size
		var voxel_y = (idx / world_size) % world_size
		var voxel_z = idx / (world_size * world_size)

		var voxel_id = Vector3i(voxel_x, voxel_y, voxel_z)
		var surface_pos = Vector3(x, y, z)

		dict[voxel_id] = surface_pos

	return dict