extends Node2D

@export var texture_size: Vector2i = Vector2i(512, 512)
@export var decay_factor: float = 0.98
@export var dissipation_factor: float = 0.12
@export var sensor_angle: float = 32.0:
	get():
		return deg_to_rad(sensor_angle)
@export var sensor_distance: float = 3.86
@export var rng_seed: int = 80085

@onready var mesh := $MeshInstance2D

## Needs to be changed from the shader as well
const AGENTS: int = 1024 * 1024

var texture: Texture2DRD
var next_texture: int = 0


func _ready():
	RenderingServer.call_on_render_thread(_initialize_compute_code.bind(texture_size))
	var mat: ShaderMaterial = mesh.material
	texture = mat.get_shader_parameter("trail_tex")


func _process(_delta: float) -> void:
	next_texture = (next_texture + 1) % 2

	if texture:
		texture.texture_rd_rid = texture_rds[next_texture]

	RenderingServer.call_on_render_thread(
		_render_process.bind(
			next_texture,
			texture_size,
			decay_factor,
			dissipation_factor,
			sensor_angle,
			sensor_distance
		)
	)


func _exit_tree() -> void:
	if texture:
		texture.texture_rd_id = RID()

	RenderingServer.call_on_render_thread(_free_compute_resources)


# Everything after this point is designed to run on our rendering thread.

var rd: RenderingDevice

var decay_shader: RID
var decay_pipeline: RID

var agent_shader: RID
var agent_pipeline: RID

# We use 3 textures:
# - One to render into
# - One that contains the last frame rendered
# - One for the frame before that
var texture_rds: Array[RID] = [RID(), RID()]
var texture_sets: Array[RID] = [RID(), RID()]

var agent_buffer: RID
var agent_sets: Array[RID] = [RID(), RID()]


func _load_shader(rd_: RenderingDevice, path: String) -> RID:
	var shader_file := load(path)
	var shader_spirv: RDShaderSPIRV = shader_file.get_spirv()
	return rd_.shader_create_from_spirv(shader_spirv)


func _create_trail_uniform(texture_rd: RID) -> RDUniform:
	var uniform := RDUniform.new()
	uniform.uniform_type = RenderingDevice.UNIFORM_TYPE_IMAGE
	uniform.binding = 0
	uniform.add_id(texture_rd)
	return uniform


func _initialize_compute_code(with_texture_size: Vector2i) -> void:
	# As this becomes part of our normal frame rendering,
	# we use our main rendering device here.
	rd = RenderingServer.get_rendering_device()

	# Create shaders
	decay_shader = _load_shader(rd, "res://slime_simulation_cs/shaders/decay_cs.glsl")
	decay_pipeline = rd.compute_pipeline_create(decay_shader)

	agent_shader = _load_shader(rd, "res://slime_simulation_cs/shaders/agents_cs.glsl")
	agent_pipeline = rd.compute_pipeline_create(agent_shader)

	# Create agent buffer
	var rng := RandomNumberGenerator.new()
	rng.seed = rng_seed
	var agent_positions := PackedVector2Array()
	var agent_angles := PackedFloat32Array()
	# This must match what is set in the shader, if fixed size
	for n in AGENTS:
		var new_position := Vector2(
			rng.randf_range(0, with_texture_size.x), rng.randf_range(0, with_texture_size.y)
		)
		agent_positions.push_back(new_position)
		agent_angles.push_back(rng.randf_range(0, TAU))

	var agent_data := PackedByteArray(agent_positions.to_byte_array())
	agent_data.append_array(agent_angles.to_byte_array())
	agent_buffer = rd.storage_buffer_create(
		agent_data.size(), agent_data, 0, RenderingDevice.BUFFER_CREATION_AS_STORAGE_BIT
	)

	# Create trail texture
	var tf: RDTextureFormat = RDTextureFormat.new()
	tf.format = RenderingDevice.DATA_FORMAT_R32_SFLOAT
	tf.texture_type = RenderingDevice.TEXTURE_TYPE_2D
	tf.width = with_texture_size.x
	tf.height = with_texture_size.y
	tf.depth = 1
	tf.array_layers = 1
	tf.mipmaps = 1
	tf.usage_bits = (
		RenderingDevice.TEXTURE_USAGE_SAMPLING_BIT
		| RenderingDevice.TEXTURE_USAGE_COLOR_ATTACHMENT_BIT
		| RenderingDevice.TEXTURE_USAGE_STORAGE_BIT
		| RenderingDevice.TEXTURE_USAGE_CAN_UPDATE_BIT
		| RenderingDevice.TEXTURE_USAGE_CAN_COPY_TO_BIT
	)

	var agent_uniforms: Array[RDUniform] = []
	for i in 2:
		texture_rds[i] = rd.texture_create(tf, RDTextureView.new(), [])
		# Make sure our textures are cleared.
		rd.texture_clear(texture_rds[i], Color.BLACK, 0, 1, 0, 1)

		var trail_uniform := _create_trail_uniform(texture_rds[i])
		texture_sets[i] = rd.uniform_set_create([trail_uniform], decay_shader, 0)

	var txtr_uniform_1 := RDUniform.new()
	txtr_uniform_1.uniform_type = RenderingDevice.UNIFORM_TYPE_IMAGE
	txtr_uniform_1.binding = 0
	txtr_uniform_1.add_id(texture_rds[0])
	agent_uniforms.append(txtr_uniform_1)

	var txtr_uniform_2 := RDUniform.new()
	txtr_uniform_2.uniform_type = RenderingDevice.UNIFORM_TYPE_IMAGE
	txtr_uniform_2.binding = 1
	txtr_uniform_2.add_id(texture_rds[1])
	agent_uniforms.append(txtr_uniform_2)

	var uniform := RDUniform.new()
	uniform.uniform_type = RenderingDevice.UNIFORM_TYPE_STORAGE_BUFFER
	uniform.binding = 2
	uniform.add_id(agent_buffer)
	agent_sets[0] = rd.uniform_set_create(
		[txtr_uniform_1, txtr_uniform_2, uniform], agent_shader, 0
	)
	agent_sets[1] = rd.uniform_set_create(
		[txtr_uniform_2, txtr_uniform_1, uniform], agent_shader, 0
	)


func _render_process(
	with_next_texture: int,
	tex_size: Vector2i,
	decay_fctr: float,
	dissipation_fctr: float,
	snsr_angle: float,
	snsr_distance: float
) -> void:
	var next_set := texture_sets[with_next_texture]
	var current_set := texture_sets[(with_next_texture - 1) % 2]

	# Run agents
	var push_constant := PackedFloat32Array()
	push_constant.push_back(tex_size.x)
	push_constant.push_back(tex_size.y)
	push_constant.push_back(snsr_angle)
	push_constant.push_back(snsr_distance)

	var agent_compute_list := rd.compute_list_begin()
	rd.compute_list_bind_compute_pipeline(agent_compute_list, agent_pipeline)
	rd.compute_list_bind_uniform_set(agent_compute_list, agent_sets[with_next_texture], 0)
	rd.compute_list_set_push_constant(
		agent_compute_list, push_constant.to_byte_array(), push_constant.size() * 4
	)
	@warning_ignore("integer_division")
	rd.compute_list_dispatch(agent_compute_list, 64, 1, 1)
	rd.compute_list_end()

	# Decay and dissipate
	push_constant.clear()
	push_constant.push_back(tex_size.x)
	push_constant.push_back(tex_size.y)
	push_constant.push_back(decay_fctr)
	push_constant.push_back(dissipation_fctr)

	@warning_ignore("integer_division")
	var x_groups := (tex_size.x - 1) / 8 + 1
	@warning_ignore("integer_division")
	var y_groups := (tex_size.y - 1) / 8 + 1

	# Run the compute shader
	var compute_list := rd.compute_list_begin()
	rd.compute_list_bind_compute_pipeline(compute_list, decay_pipeline)
	rd.compute_list_bind_uniform_set(compute_list, current_set, 0)
	rd.compute_list_bind_uniform_set(compute_list, next_set, 1)
	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, 1)
	rd.compute_list_end()


func _free_compute_resources() -> void:
	# Note that our sets and pipeline are cleaned up automatically
	# as they are dependencies :P
	for i in 2:
		if texture_rds[i]:
			rd.free_rid(texture_rds[i])

	if decay_shader:
		rd.free_rid(decay_shader)
	if agent_shader:
		rd.free_rid(agent_shader)