package gpu_tests

import bk "../backend"
import compiler "../compiler"
import "core:mem"
import "core:testing"

Recording_Call_Kind :: enum {
	Begin_Frame,
	End_Frame,
	Begin_Render_Pass,
	Begin_Default_Pass,
	End_Render_Pass,
	Set_Viewport,
	Set_Scissor,
	Create_Buffer,
	Destroy_Buffer,
	Create_Graphics_Pipeline,
	Destroy_Graphics_Pipeline,
	Bind_Pipeline,
	Bind_Descriptor,
	Push_Constants,
	Bind_Vertex_Buffer,
	Bind_Index_Buffer,
	Draw,
	Draw_Indexed,
	Draw_Indirect,
	Draw_Indexed_Indirect,
	Create_Compute_Pipeline,
	Destroy_Compute_Pipeline,
	Bind_Compute_Pipeline,
	Dispatch_Compute,
	Compute_Barrier,
	Create_Shader,
	Destroy_Shader,
}

Recording_Call :: struct {
	kind: Recording_Call_Kind,

	frame_index: u32,

	x_i: i32,
	y_i: i32,
	width_u: u32,
	height_u: u32,
	x_f: f32,
	y_f: f32,
	width_f: f32,
	height_f: f32,

	pipeline: bk.Pipeline_Handle,
	descriptor: bk.Descriptor_Handle,
	descriptor_index: u32,
	buffer: bk.Buffer_Handle,
	buffer_slot: u32,
	buffer_offset: u64,
	buffer_stride: u32,
	argument_buffer: bk.Buffer_Handle,
	argument_offset: u64,
	draw_count: u32,
	stride: u32,
	color_count: u32,
	color_formats: [bk.MAX_COLOR_TARGETS]bk.Format,
	color_write_masks: [bk.MAX_COLOR_TARGETS]u8,
	stencil: bk.Stencil_State,
	color_views: [bk.MAX_COLOR_TARGETS]bk.Texture_Handle,
	clear_colors: [bk.MAX_COLOR_TARGETS][4]f32,
	vertex_binding_count: u32,
	vertex_bindings: [16]bk.Vertex_Binding,
	vertex_attribute_count: u32,
	vertex_attributes: [16]bk.Vertex_Attribute,

	push_constant_size: u32,
	push_constant_bytes: [compiler.IMPORTED_PUSH_CONSTANT_CAP]u8,

	vertex_count: u32,
	index_count: u32,
	instance_count: u32,
	first_vertex: u32,
	first_index: u32,
	vertex_offset: i32,
	first_instance: u32,
}

Recording_Buffer :: struct {
	handle: bk.Buffer_Handle,
	data: []u8,
	mapped: bool,
}

Recording_Backend_State :: struct {
	calls: [dynamic]Recording_Call,
	buffers: [dynamic]Recording_Buffer,

	shader_desc: bk.Shader_Module_Desc,
	extent: bk.Extent,

	next_buffer: u64,
	next_pipeline: u64,
	next_shader: u64,
	current_frame_index: u32,
	next_frame_index: u32,

	begin_pass_count: int,
	end_pass_count: int,
	draw_count: int,
	vertex_count: u32,
	indexed_count: u32,
	instance_count: u32,
	bound_vertex_buffer: bk.Buffer_Handle,
	bound_index_buffer: bk.Buffer_Handle,
	bound_pipeline: bk.Pipeline_Handle,
	bound_descriptor: bk.Descriptor_Handle,
	push_constant_size: u32,
}

@(thread_local) recording_active: ^Recording_Backend_State

recording_backend_init :: proc() -> Recording_Backend_State {
	return {
		calls = make([dynamic]Recording_Call, 0, 128),
		buffers = make([dynamic]Recording_Buffer, 0, 16),
		extent = {width = 640, height = 480},
		next_buffer = 12,
		next_pipeline = 31,
		next_shader = 77,
	}
}

recording_backend_destroy :: proc(rec: ^Recording_Backend_State) {
	if rec == nil {
		return
	}
	for &buffer in rec.buffers {
		if buffer.data != nil {
			delete(buffer.data)
		}
	}
	delete(rec.calls)
	delete(rec.buffers)
	rec^ = {}
	if recording_active == rec {
		recording_active = nil
	}
}

recording_backend_reset :: proc(rec: ^Recording_Backend_State) {
	if rec == nil {
		return
	}
	clear(&rec.calls)
	rec.begin_pass_count = 0
	rec.end_pass_count = 0
	rec.draw_count = 0
	rec.vertex_count = 0
	rec.indexed_count = 0
	rec.instance_count = 0
	rec.bound_vertex_buffer = bk.NULL_BUFFER
	rec.bound_index_buffer = bk.NULL_BUFFER
	rec.bound_pipeline = bk.NULL_PIPELINE
	rec.bound_descriptor = bk.NULL_DESCRIPTOR
	rec.push_constant_size = 0
	rec.shader_desc = {}
}

recording_backend_use :: proc(rec: ^Recording_Backend_State) -> bk.Backend {
	recording_active = rec
	return {
		begin_frame = recording_begin_frame,
		end_frame = recording_end_frame,
		get_extent = recording_get_extent,
		current_frame_index = recording_current_frame_index,

		begin_render_pass = recording_begin_render_pass,
		begin_default_pass = recording_begin_default_pass,
		end_render_pass = recording_end_render_pass,
		set_viewport = recording_set_viewport,
		set_scissor = recording_set_scissor,

		create_buffer = recording_create_buffer,
		create_buffer_staged = recording_create_buffer_staged,
		destroy_buffer = recording_destroy_buffer,
		map_buffer = recording_map_buffer,
		unmap_buffer = recording_unmap_buffer,
		get_buffer_mapped = recording_get_buffer_mapped,

		create_graphics_pipeline = recording_create_graphics_pipeline,
		destroy_graphics_pipeline = recording_destroy_graphics_pipeline,
		bind_graphics_pipeline = recording_bind_graphics_pipeline,
		push_constants = recording_push_constants,
		bind_descriptor_set = recording_bind_descriptor_set,
		bind_vertex_buffer = recording_bind_vertex_buffer,
		bind_vertex_buffer_slot = recording_bind_vertex_buffer_slot,
		bind_index_buffer = recording_bind_index_buffer,
		draw = recording_draw,
		draw_indexed = recording_draw_indexed,
		draw_indirect = recording_draw_indirect,
		draw_indexed_indirect = recording_draw_indexed_indirect,
		create_compute_pipeline = recording_create_compute_pipeline,
		destroy_compute_pipeline = recording_destroy_compute_pipeline,
		bind_compute_pipeline = recording_bind_compute_pipeline,
		dispatch_compute = recording_dispatch_compute,
		compute_barrier = recording_compute_barrier,

		create_shader_module = recording_create_shader_module,
		destroy_shader = recording_destroy_shader,

		destroy_render_pass = recording_destroy_render_pass,
		destroy_framebuffer = recording_destroy_framebuffer,
	}
}

recording_expect_kinds :: proc(
	t: ^testing.T,
	rec: ^Recording_Backend_State,
	expected: []Recording_Call_Kind,
) {
	testing.expect_value(t, len(rec.calls), len(expected))
	if len(rec.calls) != len(expected) {
		return
	}
	for kind, i in expected {
		testing.expect_value(t, rec.calls[i].kind, kind)
	}
}

recording_record :: proc(call: Recording_Call) {
	if recording_active == nil {
		return
	}
	append(&recording_active.calls, call)
}

recording_begin_frame :: proc(clear_color: [4]f32) -> (bk.Frame_Context, bool) {
	_ = clear_color
	rec := recording_active
	if rec == nil {
		return {}, false
	}
	rec.current_frame_index = rec.next_frame_index
	ctx := bk.Frame_Context {frame_index = rec.current_frame_index}
	recording_record({kind = .Begin_Frame, frame_index = ctx.frame_index})
	return ctx, true
}

recording_end_frame :: proc(ctx: bk.Frame_Context) -> bool {
	rec := recording_active
	if rec != nil {
		rec.next_frame_index = (ctx.frame_index + 1) % bk.MAX_FRAMES_IN_FLIGHT
	}
	recording_record({kind = .End_Frame, frame_index = ctx.frame_index})
	return false
}

recording_current_frame_index :: proc() -> u32 {
	if recording_active == nil {
		return 0
	}
	return recording_active.current_frame_index
}

recording_begin_render_pass :: proc(ctx: bk.Frame_Context, desc: bk.Render_Pass_Begin_Desc) {
	if recording_active != nil {
		recording_active.begin_pass_count += 1
	}
	recording_record({
		kind = .Begin_Render_Pass,
		frame_index = ctx.frame_index,
		color_count = desc.color_count,
		clear_colors = desc.clear_colors,
	})
}

recording_begin_default_pass :: proc(ctx: bk.Frame_Context, clear_color: [4]f32) {
	_ = clear_color
	if recording_active != nil {
		recording_active.begin_pass_count += 1
	}
	recording_record({kind = .Begin_Default_Pass, frame_index = ctx.frame_index})
}

recording_end_render_pass :: proc(ctx: bk.Frame_Context) {
	if recording_active != nil {
		recording_active.end_pass_count += 1
	}
	recording_record({kind = .End_Render_Pass, frame_index = ctx.frame_index})
}

recording_set_viewport :: proc(ctx: bk.Frame_Context, x, y, w, h: f32) {
	recording_record({kind = .Set_Viewport, frame_index = ctx.frame_index, x_f = x, y_f = y, width_f = w, height_f = h})
}

recording_set_scissor :: proc(ctx: bk.Frame_Context, x, y: i32, w, h: u32) {
	recording_record({kind = .Set_Scissor, frame_index = ctx.frame_index, x_i = x, y_i = y, width_u = w, height_u = h})
}

recording_get_extent :: proc() -> bk.Extent {
	if recording_active == nil {
		return {width = 640, height = 480}
	}
	return recording_active.extent
}

recording_find_buffer :: proc(
	rec: ^Recording_Backend_State,
	handle: bk.Buffer_Handle,
) -> (^Recording_Buffer, bool) {
	if rec == nil {
		return nil, false
	}
	for &buffer in rec.buffers {
		if buffer.handle == handle {
			return &buffer, true
		}
	}
	return nil, false
}

recording_create_buffer :: proc(desc: bk.Buffer_Desc) -> (bk.Buffer_Handle, bool) {
	rec := recording_active
	if rec == nil {
		return bk.NULL_BUFFER, false
	}
	size := int(desc.size)
	if size <= 0 {
		size = 1
	}
	handle := bk.Buffer_Handle(rec.next_buffer)
	rec.next_buffer += 1
	append(&rec.buffers, Recording_Buffer {handle = handle, data = make([]u8, size)})
	recording_record({kind = .Create_Buffer, buffer = handle})
	return handle, true
}

recording_create_buffer_staged :: proc(
	data: rawptr,
	size: int,
	usage: bk.Buffer_Usage_Flags,
) -> (
	bk.Buffer_Handle,
	bool,
) {
	_ = usage
	handle, ok := recording_create_buffer({size = u64(size)})
	if ok && data != nil {
		if buffer, buffer_ok := recording_find_buffer(recording_active, handle); buffer_ok {
			mem.copy(raw_data(buffer.data), data, size)
		}
	}
	return handle, ok
}

recording_destroy_buffer :: proc(handle: bk.Buffer_Handle) {
	rec := recording_active
	if rec != nil {
		for &buffer, i in rec.buffers {
			if buffer.handle == handle {
				if buffer.data != nil {
					delete(buffer.data)
				}
				unordered_remove(&rec.buffers, i)
				break
			}
		}
	}
	recording_record({kind = .Destroy_Buffer, buffer = handle})
}

recording_map_buffer :: proc(handle: bk.Buffer_Handle) -> rawptr {
	buffer, ok := recording_find_buffer(recording_active, handle)
	if !ok {
		return nil
	}
	buffer.mapped = true
	return raw_data(buffer.data)
}

recording_unmap_buffer :: proc(handle: bk.Buffer_Handle) {
	if buffer, ok := recording_find_buffer(recording_active, handle); ok {
		buffer.mapped = false
	}
}

recording_get_buffer_mapped :: proc(handle: bk.Buffer_Handle) -> rawptr {
	buffer, ok := recording_find_buffer(recording_active, handle)
	if !ok || !buffer.mapped {
		return nil
	}
	return raw_data(buffer.data)
}

recording_create_graphics_pipeline :: proc(desc: bk.Pipeline_Desc) -> (bk.Pipeline_Handle, bool) {
	rec := recording_active
	if rec == nil {
		return bk.NULL_PIPELINE, false
	}
	handle := bk.Pipeline_Handle(rec.next_pipeline)
	rec.next_pipeline += 1
	recording_record({
		kind = .Create_Graphics_Pipeline,
		pipeline = handle,
		color_count = desc.color_attachment_count,
		color_formats = desc.color_formats,
		color_write_masks = desc.color_write_masks,
		stencil = desc.stencil,
		vertex_binding_count = u32(len(desc.vertex_bindings)),
		vertex_attribute_count = u32(len(desc.vertex_attributes)),
	})
	call := &rec.calls[len(rec.calls) - 1]
	for binding, i in desc.vertex_bindings {
		if i >= len(call.vertex_bindings) {break}
		call.vertex_bindings[i] = binding
	}
	for attr, i in desc.vertex_attributes {
		if i >= len(call.vertex_attributes) {break}
		call.vertex_attributes[i] = attr
	}
	return handle, true
}

recording_destroy_graphics_pipeline :: proc(handle: bk.Pipeline_Handle) {
	recording_record({kind = .Destroy_Graphics_Pipeline, pipeline = handle})
}

recording_bind_graphics_pipeline :: proc(ctx: bk.Frame_Context, handle: bk.Pipeline_Handle) {
	if recording_active != nil {
		recording_active.bound_pipeline = handle
	}
	recording_record({kind = .Bind_Pipeline, frame_index = ctx.frame_index, pipeline = handle})
}

recording_bind_descriptor_set :: proc(
	ctx: bk.Frame_Context,
	pipeline: bk.Pipeline_Handle,
	set: bk.Descriptor_Handle,
	index: u32,
) {
	if recording_active != nil {
		recording_active.bound_descriptor = set
	}
	recording_record({
		kind = .Bind_Descriptor,
		frame_index = ctx.frame_index,
		pipeline = pipeline,
		descriptor = set,
		descriptor_index = index,
	})
}

recording_push_constants :: proc(
	ctx: bk.Frame_Context,
	pipeline: bk.Pipeline_Handle,
	stages: bk.Shader_Stage_Flags,
	offset, size: u32,
	data: rawptr,
) {
	_ = stages
	_ = offset
	call := Recording_Call {
		kind = .Push_Constants,
		frame_index = ctx.frame_index,
		pipeline = pipeline,
		push_constant_size = size,
	}
	if data != nil && size <= compiler.IMPORTED_PUSH_CONSTANT_CAP {
		mem.copy(&call.push_constant_bytes[0], data, int(size))
	}
	if recording_active != nil {
		recording_active.push_constant_size = size
	}
	recording_record(call)
}

recording_bind_vertex_buffer :: proc(ctx: bk.Frame_Context, handle: bk.Buffer_Handle) {
	recording_bind_vertex_buffer_slot(ctx, 0, handle, 0, 0)
}

recording_bind_vertex_buffer_slot :: proc(
	ctx: bk.Frame_Context,
	slot: u32,
	handle: bk.Buffer_Handle,
	offset: u64,
	stride: u32,
) {
	if recording_active != nil {
		if slot == 0 {
			recording_active.bound_vertex_buffer = handle
		}
	}
	recording_record({
		kind = .Bind_Vertex_Buffer,
		frame_index = ctx.frame_index,
		buffer = handle,
		buffer_slot = slot,
		buffer_offset = offset,
		buffer_stride = stride,
	})
}

recording_bind_index_buffer :: proc(ctx: bk.Frame_Context, handle: bk.Buffer_Handle) {
	if recording_active != nil {
		recording_active.bound_index_buffer = handle
	}
	recording_record({kind = .Bind_Index_Buffer, frame_index = ctx.frame_index, buffer = handle})
}

recording_draw :: proc(
	ctx: bk.Frame_Context,
	vertex_count, instance_count: u32,
	first_vertex: u32,
	first_instance: u32,
) {
	if recording_active != nil {
		recording_active.draw_count += 1
		recording_active.vertex_count = vertex_count
		recording_active.instance_count = instance_count
	}
	recording_record({
		kind = .Draw,
		frame_index = ctx.frame_index,
		vertex_count = vertex_count,
		instance_count = instance_count,
		first_vertex = first_vertex,
		first_instance = first_instance,
	})
}

recording_draw_indexed :: proc(
	ctx: bk.Frame_Context,
	index_count, instance_count: u32,
	first_index: u32,
	vertex_offset: i32,
	first_instance: u32,
) {
	if recording_active != nil {
		recording_active.draw_count += 1
		recording_active.indexed_count = index_count
		recording_active.instance_count = instance_count
	}
	recording_record({
		kind = .Draw_Indexed,
		frame_index = ctx.frame_index,
		index_count = index_count,
		instance_count = instance_count,
		first_index = first_index,
		vertex_offset = vertex_offset,
		first_instance = first_instance,
	})
}

recording_draw_indirect :: proc(
	ctx: bk.Frame_Context,
	argument_buffer: bk.Buffer_Handle,
	argument_offset: u64,
	draw_count: u32,
	stride: u32,
) {
	if recording_active != nil {
		recording_active.draw_count += int(draw_count)
	}
	recording_record({
		kind = .Draw_Indirect,
		frame_index = ctx.frame_index,
		argument_buffer = argument_buffer,
		argument_offset = argument_offset,
		draw_count = draw_count,
		stride = stride,
	})
}

recording_draw_indexed_indirect :: proc(
	ctx: bk.Frame_Context,
	argument_buffer: bk.Buffer_Handle,
	argument_offset: u64,
	draw_count: u32,
	stride: u32,
) {
	if recording_active != nil {
		recording_active.draw_count += int(draw_count)
	}
	recording_record({
		kind = .Draw_Indexed_Indirect,
		frame_index = ctx.frame_index,
		argument_buffer = argument_buffer,
		argument_offset = argument_offset,
		draw_count = draw_count,
		stride = stride,
	})
}

recording_create_compute_pipeline :: proc(
	shader: bk.Shader_Handle,
	num_buffers: u32,
	push_constant_size: u32,
) -> (
	bk.Pipeline_Handle,
	bool,
) {
	_ = shader
	_ = num_buffers
	_ = push_constant_size
	rec := recording_active
	if rec == nil {
		return bk.NULL_PIPELINE, false
	}
	handle := bk.Pipeline_Handle(rec.next_pipeline)
	rec.next_pipeline += 1
	recording_record({kind = .Create_Compute_Pipeline, pipeline = handle})
	return handle, true
}

recording_destroy_compute_pipeline :: proc(handle: bk.Pipeline_Handle) {
	recording_record({kind = .Destroy_Compute_Pipeline, pipeline = handle})
}

recording_bind_compute_pipeline :: proc(ctx: bk.Frame_Context, handle: bk.Pipeline_Handle) {
	if recording_active != nil {
		recording_active.bound_pipeline = handle
	}
	recording_record({kind = .Bind_Compute_Pipeline, frame_index = ctx.frame_index, pipeline = handle})
}

recording_dispatch_compute :: proc(ctx: bk.Frame_Context, groups_x, groups_y, groups_z: u32) {
	recording_record({
		kind = .Dispatch_Compute,
		frame_index = ctx.frame_index,
		width_u = groups_x,
		height_u = groups_y,
		vertex_count = groups_z,
	})
}

recording_compute_barrier :: proc(ctx: bk.Frame_Context) {
	recording_record({kind = .Compute_Barrier, frame_index = ctx.frame_index})
}

recording_create_shader_module :: proc(desc: bk.Shader_Module_Desc) -> (bk.Shader_Handle, bool) {
	rec := recording_active
	if rec == nil {
		return bk.NULL_SHADER, false
	}
	rec.shader_desc = desc
	handle := bk.Shader_Handle(rec.next_shader)
	rec.next_shader += 1
	recording_record({kind = .Create_Shader})
	return handle, true
}

recording_destroy_shader :: proc(handle: bk.Shader_Handle) {
	_ = handle
	recording_record({kind = .Destroy_Shader})
}

recording_destroy_render_pass :: proc(handle: bk.Render_Pass_Handle) {
	_ = handle
}

recording_destroy_framebuffer :: proc(handle: bk.Framebuffer_Handle) {
	_ = handle
}