package renderer

import "core:log"
import bk "../backend"
import ir "../render_ir"
import "../profile"
import "../resource"

Runtime_Storage_Binding :: struct {
	binding: u32,
	storage_buffer_id: u32,
	size: u64,
}

Runtime_Compute_Dispatch :: struct {
	shader_id: u32,
	groups: [3]u32,
	storage_bindings: [ir.MAX_COMPUTE_BINDINGS]Runtime_Storage_Binding,
	storage_binding_count: u8,
	push_constants: []u8,
	barrier_after: bool,
}

@(private)
reset_frame_ir_compute :: proc(state: ^Renderer_State) -> (ir.Pass_Handle, ir.Pipeline_Handle) {
	ir.reset_frame_ir(&state.frame_ir)
	pass := ir.add_pass(&state.frame_ir, {kind = .Compute, name = "compute"})
	pipeline := ir.add_pipeline(&state.frame_ir, {kind = .Compute, name = "compute"})
	return pass, pipeline
}

execute_compute_dispatch :: proc(state: ^Renderer_State, res_state: ^resource.Resource_State, dispatch: ^Runtime_Compute_Dispatch) {
	if !state.frame_active {
		return
	}

	if state.render_pass_begun {
		log.error("gpu: dispatch_compute must be called before any draw operations in the frame")
		return
	}

	cs, cs_ok := resource.get_compute_shader(res_state, dispatch.shader_id)
	if !cs_ok {
		return
	}

	b := state.backend
	if b == nil {
		log.error("gpu/renderer: compute dispatch missing backend")
		return
	}

	ctx := state.frame_ctx
	desc_set := cs.desc_sets[ctx.frame_index]

	count := min(int(dispatch.storage_binding_count), ir.MAX_COMPUTE_BINDINGS)
	if count > 0 && desc_set == bk.NULL_DESCRIPTOR {
		log.error("gpu: compute shader has no storage-buffer descriptor set")
		return
	}
	for i in 0..<count {
		binding := dispatch.storage_bindings[i]
		sb, sb_ok := resource.get_storage_buffer(res_state, binding.storage_buffer_id)
		if !sb_ok {
			log.errorf("gpu: invalid storage buffer at binding %d", i)
			return
		}
		descriptor_size := u64(sb.size)
		if binding.size > 0 {
			if binding.size > descriptor_size {
				log.errorf("gpu: storage buffer binding %d exceeds imported buffer size", binding.binding)
				return
			}
			descriptor_size = binding.size
		}
		b.update_descriptor_buffer(
			desc_set,
			binding.binding,
			sb.buffer,
			descriptor_size,
		)
		state.profile.descriptor_updates += 1
	}

	b.bind_compute_pipeline(ctx, cs.pipeline)
	state.profile.pipeline_binds += 1
	if desc_set != bk.NULL_DESCRIPTOR {
		b.bind_descriptor_set(ctx, cs.pipeline, desc_set, 0)
		state.profile.descriptor_binds += 1
	}

	if len(dispatch.push_constants) > 0 {
		b.push_constants(ctx, cs.pipeline, {.Compute}, 0, u32(len(dispatch.push_constants)), raw_data(dispatch.push_constants))
		state.profile.push_constants += 1
	}

	b.dispatch_compute(ctx, dispatch.groups[0], dispatch.groups[1], dispatch.groups[2])
	state.profile.dispatch_calls += 1

	if dispatch.barrier_after {
		b.compute_barrier(ctx)
	}
}

flush_frame_ir_compute :: proc(state: ^Renderer_State, res_state: ^resource.Resource_State) {
	if ir.command_count(&state.frame_ir) == 0 {
		return
	}

	prof_start := profile.now()

	for _, i in state.frame_ir.commands {
		cmd := ir.Command_Handle(i + 1)
		dispatch, ok := ir.get_dispatch(&state.frame_ir, cmd)
		if !ok {
			continue
		}
		state.profile.compute_dispatches += 1
		runtime_dispatch: Runtime_Compute_Dispatch = {
			shader_id = dispatch.shader_id,
			groups = dispatch.groups,
			push_constants = dispatch.push_constants,
			barrier_after = dispatch.barrier_after,
		}
		execute_compute_dispatch(state, res_state, &runtime_dispatch)
	}

	ir.reset_frame_ir(&state.frame_ir)
	state.profile.compute_lower_ns += profile.elapsed_ns(prof_start, profile.now())
}

lower_compute_dispatch_ir :: proc(
	state: ^Renderer_State,
	res_state: ^resource.Resource_State,
	shader_id: u32,
	groups_x, groups_y, groups_z: u32,
	push_constants: rawptr,
	push_constants_size: u32,
	storage_buffer_ids: []u32,
) {
	push_data: []u8
	if push_constants != nil && push_constants_size > 0 {
		push_data = (cast([^]u8)push_constants)[:int(push_constants_size)]
	}
	dispatch := Runtime_Compute_Dispatch{
		shader_id = shader_id,
		groups = {groups_x, groups_y, groups_z},
		push_constants = push_data,
		barrier_after = true,
	}
	count := min(len(storage_buffer_ids), ir.MAX_COMPUTE_BINDINGS)
	dispatch.storage_binding_count = u8(count)
	for i in 0..<count {
		dispatch.storage_bindings[i] = {
			binding = u32(i),
			storage_buffer_id = storage_buffer_ids[i],
		}
	}
	state.profile.compute_dispatches += 1
	execute_compute_dispatch(state, res_state, &dispatch)
}