package render_ir

import "core:mem"

Resource_Access_Kind :: enum {
	Color_Attachment_Write,
	Depth_Attachment_Write,
	Sampled_Texture_Read,
	Storage_Buffer_Read_Write,
	Vertex_Buffer_Read,
	Index_Buffer_Read,
	Instance_Buffer_Read,
	Indirect_Argument_Read,
	Access_Metadata_Missing,
}

Resource_Access :: struct {
	command:  Command_Handle,
	pass:     Pass_Handle,
	resource: Resource_Handle,
	kind:     Resource_Access_Kind,
}

Resource_Access_List :: [dynamic]Resource_Access

destroy_resource_accesses :: proc(accesses: ^Resource_Access_List) {
	delete(accesses^)
	accesses^ = nil
}

collect_resource_accesses :: proc(
	frame: ^Frame_IR,
	allocator: mem.Allocator = context.allocator,
) -> Resource_Access_List {
	accesses := make(Resource_Access_List, 0, DEFAULT_COMMAND_CAP, allocator)
	if frame == nil {
		return accesses
	}
	plan := plan_frame_commands(frame, context.temp_allocator)
	defer command_plan_destroy(&plan)
	for command in plan {
		if command.culled {
			continue
		}
		switch command.record.kind {
		case .Draw:
			draw, draw_ok := get_draw(frame, command.handle)
			if !draw_ok {continue}
			collect_draw_accesses(frame, command.handle, draw, &accesses)
		case .Dispatch:
			dispatch, dispatch_ok := get_dispatch(frame, command.handle)
			if !dispatch_ok {continue}
			collect_dispatch_accesses(command.handle, dispatch, &accesses)
		}
	}
	return accesses
}

validate_resource_accesses :: proc(frame: ^Frame_IR, diagnostics: ^Diagnostic_List) -> bool {
	if frame == nil || diagnostics == nil {
		return false
	}
	ok := true
	plan := plan_frame_commands(frame, context.temp_allocator)
	defer command_plan_destroy(&plan)
	written_targets := make([dynamic]Resource_Handle, 0, DEFAULT_PASS_CAP, context.temp_allocator)
	known_frame_targets := collect_frame_target_resources(frame, context.temp_allocator)
	defer delete(known_frame_targets)
	seen_present := false
	for command in plan {
		if command.culled {
			continue
		}
		switch command.record.kind {
		case .Draw:
			draw, draw_ok := get_draw(frame, command.handle)
			if !draw_ok {continue}
			if draw_targets_present(frame, draw) {
				seen_present = true
			} else if seen_present && draw_targets_offscreen(frame, draw) {
				add_command_error(
					diagnostics,
					.Invalid_Order,
					command.handle,
					"offscreen work after present pass is unsupported",
				)
				ok = false
			}
			pass := effective_draw_pass(frame, draw)
			targets := pass_target_resources(frame, pass, context.temp_allocator)
			sampled := draw_sampled_resources(frame, draw, context.temp_allocator)
			for target in targets {
				for read in sampled {
					if target.resource != INVALID_RESOURCE && target.resource == read {
						add_command_error(
							diagnostics,
							.Resource_Feedback,
							command.handle,
							"draw samples a resource currently bound as a render target",
						)
						ok = false
					}
				}
				if target.resource != INVALID_RESOURCE {
					append_unique_resource(&written_targets, target.resource)
				}
			}
			for read in sampled {
				if resource_in(read, known_frame_targets[:]) && !resource_in(read, written_targets[:]) {
					add_command_error(
						diagnostics,
						.Invalid_Order,
						command.handle,
						"draw samples a frame target before it is written in planned order",
					)
					ok = false
				}
			}
			delete(targets)
			delete(sampled)
		case .Dispatch:
		}
	}
	return ok
}

@(private)
collect_draw_accesses :: proc(
	frame: ^Frame_IR,
	command: Command_Handle,
	draw: ^Draw_Command,
	accesses: ^Resource_Access_List,
) {
	pass := effective_draw_pass(frame, draw)
	append_pass_target_accesses(frame, command, pass, accesses)
	geometry := draw.packet.geometry
	if geometry.resource != INVALID_RESOURCE {
		kind := Resource_Access_Kind.Index_Buffer_Read if geometry.kind == .Indexed_Mesh else Resource_Access_Kind.Vertex_Buffer_Read
		append(accesses, Resource_Access{command = command, pass = pass, resource = geometry.resource, kind = kind})
	}
	if draw.packet.instances.resource != INVALID_RESOURCE {
		append(
			accesses,
			Resource_Access {
				command = command,
				pass = pass,
				resource = draw.packet.instances.resource,
				kind = .Instance_Buffer_Read,
			},
		)
	}
	if draw.packet.instances.indirect != INVALID_RESOURCE {
		append(
			accesses,
			Resource_Access {
				command = command,
				pass = pass,
				resource = draw.packet.instances.indirect,
				kind = .Indirect_Argument_Read,
			},
		)
	}
	if material, material_ok := get_material(frame, draw.packet.material); material_ok {
		if material.texture_id != 0 || material.normal_map_texture_id != 0 {
			append(accesses, Resource_Access{command = command, pass = pass, kind = .Access_Metadata_Missing})
		}
		append_descriptor_sampled_accesses(frame, command, pass, material.descriptor_set, accesses)
	}
}

@(private)
collect_dispatch_accesses :: proc(
	command: Command_Handle,
	dispatch: ^Dispatch_Command,
	accesses: ^Resource_Access_List,
) {
	for i in 0..<int(dispatch.storage_binding_count) {
		binding := dispatch.storage_bindings[i]
		if binding.resource == INVALID_RESOURCE {continue}
		append(
			accesses,
			Resource_Access {
				command = command,
				pass = dispatch.pass,
				resource = binding.resource,
				kind = .Storage_Buffer_Read_Write,
			},
		)
	}
}

@(private)
append_pass_target_accesses :: proc(
	frame: ^Frame_IR,
	command: Command_Handle,
	pass: Pass_Handle,
	accesses: ^Resource_Access_List,
) {
	targets := pass_target_resources(frame, pass, context.temp_allocator)
	defer delete(targets)
	for target in targets {
		append(accesses, Resource_Access{command = command, pass = pass, resource = target.resource, kind = target.kind})
	}
}

@(private)
append_descriptor_sampled_accesses :: proc(
	frame: ^Frame_IR,
	command: Command_Handle,
	pass: Pass_Handle,
	set_handle: Descriptor_Set_Handle,
	accesses: ^Resource_Access_List,
) {
	if set_handle == INVALID_DESCRIPTOR_SET {return}
	set, set_ok := get_descriptor_set(frame, set_handle)
	if !set_ok {return}
	for i in 0..<int(set.binding_count) {
		binding := set.bindings[i]
		if binding.type != .Combined_Image_Sampler || binding.resource == INVALID_RESOURCE {
			continue
		}
		append(
			accesses,
			Resource_Access {
				command = command,
				pass = pass,
				resource = binding.resource,
				kind = .Sampled_Texture_Read,
			},
		)
	}
}

@(private)
Target_Access :: struct {
	resource: Resource_Handle,
	kind:     Resource_Access_Kind,
}

@(private)
pass_target_resources :: proc(
	frame: ^Frame_IR,
	pass_handle: Pass_Handle,
	allocator: mem.Allocator,
) -> [dynamic]Target_Access {
	targets := make([dynamic]Target_Access, 0, MAX_RENDER_TARGETS + 1, allocator)
	pass, pass_ok := get_pass(frame, pass_handle)
	if !pass_ok || pass.kind != .Render {
		return targets
	}
	for i in 0..<int(pass.color_target_count) {
		resource := pass.color_targets[i].resource
		if resource != INVALID_RESOURCE {
			append(&targets, Target_Access{resource = resource, kind = .Color_Attachment_Write})
		}
	}
	if pass.has_depth_target && pass.depth_target.resource != INVALID_RESOURCE {
		append(&targets, Target_Access{resource = pass.depth_target.resource, kind = .Depth_Attachment_Write})
	}
	return targets
}

@(private)
draw_sampled_resources :: proc(
	frame: ^Frame_IR,
	draw: ^Draw_Command,
	allocator: mem.Allocator,
) -> [dynamic]Resource_Handle {
	reads := make([dynamic]Resource_Handle, 0, MAX_DESCRIPTOR_BINDINGS, allocator)
	if draw == nil {return reads}
	if material, material_ok := get_material(frame, draw.packet.material); material_ok {
		if set, set_ok := get_descriptor_set(frame, material.descriptor_set); set_ok {
			for i in 0..<int(set.binding_count) {
				binding := set.bindings[i]
				if binding.type == .Combined_Image_Sampler && binding.resource != INVALID_RESOURCE {
					append_unique_resource(&reads, binding.resource)
				}
			}
		}
	}
	return reads
}

@(private)
collect_frame_target_resources :: proc(
	frame: ^Frame_IR,
	allocator: mem.Allocator,
) -> [dynamic]Resource_Handle {
	resources := make([dynamic]Resource_Handle, 0, len(frame.passes), allocator)
	for pass in frame.passes {
		if pass.kind != .Render {continue}
		for i in 0..<int(pass.color_target_count) {
			append_unique_resource(&resources, pass.color_targets[i].resource)
		}
		if pass.has_depth_target {
			append_unique_resource(&resources, pass.depth_target.resource)
		}
	}
	return resources
}

@(private)
effective_draw_pass :: proc(frame: ^Frame_IR, draw: ^Draw_Command) -> Pass_Handle {
	if draw == nil {return INVALID_PASS}
	if draw.pass != INVALID_PASS {
		return draw.pass
	}
	if target, target_ok := get_target(frame, draw.packet.target); target_ok {
		return target.pass
	}
	return INVALID_PASS
}

@(private)
draw_targets_present :: proc(frame: ^Frame_IR, draw: ^Draw_Command) -> bool {
	if draw == nil {return false}
	if target, target_ok := get_target(frame, draw.packet.target); target_ok {
		return target.kind == .Present
	}
	return draw.pass == INVALID_PASS
}

@(private)
draw_targets_offscreen :: proc(frame: ^Frame_IR, draw: ^Draw_Command) -> bool {
	if draw == nil {return false}
	if target, target_ok := get_target(frame, draw.packet.target); target_ok {
		return target.kind == .Offscreen
	}
	return draw.pass != INVALID_PASS
}

@(private)
append_unique_resource :: proc(resources: ^[dynamic]Resource_Handle, resource: Resource_Handle) {
	if resource == INVALID_RESOURCE || resource_in(resource, resources[:]) {
		return
	}
	append(resources, resource)
}

@(private)
resource_in :: proc(resource: Resource_Handle, resources: []Resource_Handle) -> bool {
	for existing in resources {
		if existing == resource {
			return true
		}
	}
	return false
}