package backend

import glsl "core:math/linalg/glsl"

// --- Opaque handle types ---

Pipeline_Handle :: distinct u64
Buffer_Handle :: distinct u64
Texture_Handle :: distinct u64
Shader_Handle :: distinct u64
Descriptor_Handle :: distinct u64
Render_Pass_Handle :: distinct u64
Framebuffer_Handle :: distinct u64
Sampler_Handle :: distinct u64

NULL_PIPELINE :: Pipeline_Handle(0)
NULL_BUFFER :: Buffer_Handle(0)
NULL_TEXTURE :: Texture_Handle(0)
NULL_SHADER :: Shader_Handle(0)
NULL_DESCRIPTOR :: Descriptor_Handle(0)
NULL_RENDER_PASS :: Render_Pass_Handle(0)
NULL_FRAMEBUFFER :: Framebuffer_Handle(0)
NULL_SAMPLER :: Sampler_Handle(0)

MAX_FRAMES_IN_FLIGHT :: 2
MAX_COLOR_TARGETS :: 8
COLOR_WRITE_MASK_ALL :: u8(0x0F)

Capability :: enum {
	Compute_Dispatch,
	Storage_Buffers,
	Indirect_Draws,
	Offscreen_Targets,
	Multiple_Render_Targets,
	Stencil_Clips,
	Sampled_Targets,
}
Capability_Flags :: bit_set[Capability]

Capabilities :: struct {
	features:               Capability_Flags,
	max_color_targets:      u32,
	max_push_constant_size: u32,
	max_frames_in_flight:   u32,
}

no_capabilities :: proc() -> Capabilities {
	return {}
}

implemented_base_capabilities :: proc(
	max_color_targets, max_push_constant_size: u32,
) -> Capabilities {
	features := Capability_Flags{.Compute_Dispatch, .Storage_Buffers, .Indirect_Draws, .Offscreen_Targets, .Sampled_Targets, .Stencil_Clips}
	if max_color_targets > 1 {
		features += {.Multiple_Render_Targets}
	}
	return {
		features = features,
		max_color_targets = max_color_targets,
		max_push_constant_size = max_push_constant_size,
		max_frames_in_flight = MAX_FRAMES_IN_FLIGHT,
	}
}

supports :: proc(caps: Capabilities, capability: Capability) -> bool {
	return capability in caps.features
}

supports_color_target_count :: proc(caps: Capabilities, count: u32) -> bool {
	return count == 0 || count <= caps.max_color_targets
}

supports_push_constant_size :: proc(caps: Capabilities, size: u32) -> bool {
	return size == 0 || size <= caps.max_push_constant_size
}

handle_index :: proc(handle: $T, max_count: int) -> (int, bool) {
	idx := int(u64(handle))
	if idx <= 0 || idx >= max_count {
		return 0, false
	}
	return idx, true
}

// --- Surface descriptors ---

Surface_Kind :: enum {
	None,
	X11,
	Wayland,
	Win32,
}

X11_Surface_Handle :: struct {
	display: rawptr,
	window:  u64,
}

Wayland_Surface_Handle :: struct {
	display: rawptr,
	surface: rawptr,
}

Win32_Surface_Handle :: struct {
	hwnd:      rawptr,
	hinstance: rawptr,
}

Surface_Desc :: struct {
	kind:    Surface_Kind,
	x11:     X11_Surface_Handle,
	wayland: Wayland_Surface_Handle,
	win32:   Win32_Surface_Handle,
}

// --- Enums ---

Buffer_Usage :: enum {
	Vertex,
	Index,
	Uniform,
	Storage,
	Indirect_Argument,
	Transfer_Src,
	Transfer_Dst,
}
Buffer_Usage_Flags :: bit_set[Buffer_Usage]

Indirect_Draw_Args :: struct {
	vertex_count:   u32,
	instance_count: u32,
	first_vertex:   u32,
	first_instance: u32,
}

Indirect_Draw_Indexed_Args :: struct {
	index_count:    u32,
	instance_count: u32,
	first_index:    u32,
	vertex_offset:  i32,
	first_instance: u32,
}

#assert(size_of(Indirect_Draw_Args) == 16)
#assert(offset_of(Indirect_Draw_Args, vertex_count) == 0)
#assert(offset_of(Indirect_Draw_Args, instance_count) == 4)
#assert(offset_of(Indirect_Draw_Args, first_vertex) == 8)
#assert(offset_of(Indirect_Draw_Args, first_instance) == 12)
#assert(size_of(Indirect_Draw_Indexed_Args) == 20)
#assert(offset_of(Indirect_Draw_Indexed_Args, index_count) == 0)
#assert(offset_of(Indirect_Draw_Indexed_Args, instance_count) == 4)
#assert(offset_of(Indirect_Draw_Indexed_Args, first_index) == 8)
#assert(offset_of(Indirect_Draw_Indexed_Args, vertex_offset) == 12)
#assert(offset_of(Indirect_Draw_Indexed_Args, first_instance) == 16)

Memory_Property :: enum {
	Device_Local,
	Host_Visible,
	Host_Coherent,
}
Memory_Property_Flags :: bit_set[Memory_Property]

Shader_Stage :: enum {
	Vertex,
	Fragment,
	Compute,
}

Shader_Format :: enum {
	SPIRV,
	HLSL,
	GLSL,
}

Shader_Module_Desc :: struct {
	stage:  Shader_Stage,
	format: Shader_Format,
	name:   string,
	entry:  string,
	data:   []u8,
}

Shader_Stage_Flag :: enum {
	Vertex,
	Fragment,
	Compute,
}
Shader_Stage_Flags :: bit_set[Shader_Stage_Flag]

Topology :: enum {
	Triangle_List,
	Triangle_Strip,
	Line_List,
	Point_List,
}

Cull_Mode :: enum {
	None,
	Front,
	Back,
	Front_And_Back,
}

Front_Face :: enum {
	Counter_Clockwise,
	Clockwise,
}

Format :: enum {
	Undefined,
	R8G8B8A8_SRGB,
	R8G8B8A8_UNORM,
	B8G8R8A8_SRGB,
	D32_SFLOAT,
	D32_SFLOAT_S8_UINT,
	D24_UNORM_S8_UINT,
}

Image_Usage :: enum {
	Sampled,
	Color_Attachment,
	Depth_Stencil_Attachment,
	Transfer_Dst,
	Transfer_Src,
}
Image_Usage_Flags :: bit_set[Image_Usage]

Image_Aspect :: enum {
	Color,
	Depth,
	Stencil,
}
Image_Aspect_Flags :: bit_set[Image_Aspect]

Descriptor_Type :: enum {
	Combined_Image_Sampler,
	Uniform_Buffer,
	Storage_Buffer,
}

Filter :: enum {
	Nearest,
	Linear,
}

Address_Mode :: enum {
	Repeat,
	Clamp_To_Edge,
	Clamp_To_Border,
}

Compare_Op :: enum {
	Never,
	Less,
	Less_Or_Equal,
	Equal,
	Greater_Or_Equal,
	Greater,
	Not_Equal,
	Always,
}

Stencil_Op :: enum {
	Keep,
	Zero,
	Replace,
	Increment_Clamp,
	Decrement_Clamp,
	Invert,
}

Stencil_Face_State :: struct {
	fail_op:       Stencil_Op,
	pass_op:       Stencil_Op,
	depth_fail_op: Stencil_Op,
	compare_op:    Compare_Op,
}

Stencil_State :: struct {
	enable:     bool,
	read_mask:  u8,
	write_mask: u8,
	reference:  u8,
	front:      Stencil_Face_State,
	back:       Stencil_Face_State,
}

Image_Layout :: enum {
	Undefined,
	Color_Attachment,
	Depth_Stencil_Attachment,
	Shader_Read_Only,
	Depth_Stencil_Read_Only,
	Transfer_Dst,
	Present_Src,
}

Attachment_Load_Op :: enum {
	Load,
	Clear,
	Dont_Care,
}

Attachment_Store_Op :: enum {
	Store,
	Dont_Care,
}

// --- Descriptor structs ---

Vertex_Attribute :: struct {
	location: u32,
	binding:  u32,
	offset:   u32,
	format:   Vertex_Format,
}

Vertex_Format :: enum {
	Float2,
	Float3,
	Float4,
}

Vertex_Binding :: struct {
	binding:    u32,
	stride:     u32,
	input_rate: Vertex_Input_Rate,
}

Vertex_Input_Rate :: enum {
	Vertex,
	Instance,
}

// --- Vertex layout descriptor ---

Vertex_Attrib :: enum u8 {
	Position,
	Normal,
	Tex_Coord,
	Tex_Coord_1,
	Color,
	Tangent,
}

MAX_VERTEX_ATTRIBS :: 8

Vertex_Layout_Entry :: struct {
	attrib: Vertex_Attrib,
	format: Vertex_Format,
	offset: u32,
}

Vertex_Layout :: struct {
	entries: [MAX_VERTEX_ATTRIBS]Vertex_Layout_Entry,
	count:   u8,
	stride:  u32,
}

Blend_Mode :: enum u8 {
	Alpha, // src*α + dst*(1-α)
	Additive, // src*α + dst
	Premultiplied_Alpha, // src + dst*(1-α)
}

Blend_Factor :: enum {
	Zero,
	One,
	Src_Alpha,
	One_Minus_Src_Alpha,
}

Blend_Factors :: struct {
	src_color: Blend_Factor,
	dst_color: Blend_Factor,
	src_alpha: Blend_Factor,
	dst_alpha: Blend_Factor,
}

blend_factors :: proc(mode: Blend_Mode) -> Blend_Factors {
	switch mode {
	case .Alpha:
		return {.Src_Alpha, .One_Minus_Src_Alpha, .One, .One_Minus_Src_Alpha}
	case .Additive:
		return {.Src_Alpha, .One, .One, .One}
	case .Premultiplied_Alpha:
		return {.One, .One_Minus_Src_Alpha, .One, .One_Minus_Src_Alpha}
	}
	return {.Src_Alpha, .One_Minus_Src_Alpha, .One, .One_Minus_Src_Alpha}
}

Pipeline_Desc :: struct {
	vert_shader:          Shader_Handle,
	frag_shader:          Shader_Handle,
	render_pass:          Render_Pass_Handle,
	topology:             Topology,
	cull_mode:            Cull_Mode,
	front_face:           Front_Face,
	enable_blending:      bool,
	blend_mode:           Blend_Mode,
	enable_depth_test:    bool,
	enable_depth_bias:    bool,
	depth_format:         Format,
	stencil:              Stencil_State,
	depth_only:           bool,
	color_attachment_count: u32,
	color_formats:        [MAX_COLOR_TARGETS]Format,
	color_write_masks:    [MAX_COLOR_TARGETS]u8,
	push_constant_size:   u32,
	push_constant_stages: Shader_Stage_Flags,
	descriptor_layouts:   []Descriptor_Handle,
	vertex_bindings:      []Vertex_Binding,
	vertex_attributes:    []Vertex_Attribute,
}

Buffer_Desc :: struct {
	size:   u64,
	usage:  Buffer_Usage_Flags,
	memory: Memory_Property_Flags,
}

Texture_Desc :: struct {
	width:  u32,
	height: u32,
	format: Format,
	usage:  Image_Usage_Flags,
}

Readback_Texture_Desc :: struct {
	texture:        Texture_Handle,
	width, height:  u32,
	current_layout: Image_Layout,
}

Sampler_Desc :: struct {
	mag_filter:     Filter,
	min_filter:     Filter,
	address_mode_u: Address_Mode,
	address_mode_v: Address_Mode,
	enable_aniso:   bool,
	enable_compare: bool,
	compare_op:     Compare_Op,
	mipmap_mode:    Filter,
}

Descriptor_Set_Layout_Binding :: struct {
	binding: u32,
	type:    Descriptor_Type,
	count:   u32,
	stages:  Shader_Stage_Flags,
}

Render_Pass_Desc :: struct {
	has_color:          bool,
	has_depth:          bool,
	color_count:        u32,
	color_formats:      [MAX_COLOR_TARGETS]Format,
	color_load_ops:     [MAX_COLOR_TARGETS]Attachment_Load_Op,
	color_store_ops:    [MAX_COLOR_TARGETS]Attachment_Store_Op,
	color_final_layouts: [MAX_COLOR_TARGETS]Image_Layout,
	color_format:       Format,
	depth_format:       Format,
	depth_only:         bool, // depth-only pass (shadows)
	color_load_op:      Attachment_Load_Op,
	color_store_op:     Attachment_Store_Op,
	color_final_layout: Image_Layout,
	depth_load_op:      Attachment_Load_Op,
	depth_store_op:     Attachment_Store_Op,
	depth_final_layout: Image_Layout,
	has_stencil:        bool,
	stencil_load_op:    Attachment_Load_Op,
	stencil_store_op:   Attachment_Store_Op,
}

Framebuffer_Desc :: struct {
	pass:       Render_Pass_Handle,
	color_count: u32,
	color_views: [MAX_COLOR_TARGETS]Texture_Handle,
	color_view: Texture_Handle,
	depth_view: Texture_Handle,
	width:      u32,
	height:     u32,
	layers:     u32,
}

Render_Pass_Begin_Desc :: struct {
	pass:        Render_Pass_Handle,
	framebuffer: Framebuffer_Handle,
	width:       u32,
	height:      u32,
	color_count: u32,
	clear_colors: [MAX_COLOR_TARGETS][4]f32,
	clear_color: [4]f32,
	clear_depth: f32,
	clear_stencil: u8,
}

// --- Frame context ---

// Opaque handle representing the current frame's recording state.
// Backends map this to their command buffer / frame index internals.
Frame_Context :: struct {
	frame_index: u32,
	_opaque:     [2]u64, // backend-specific data (e.g. command buffer handle)
}

// --- Surface extent ---

Extent :: struct {
	width:  u32,
	height: u32,
}

// --- Backend vtable ---

// Backend is a struct of proc pointers that abstract GPU operations.
// Each backend implementation (Vulkan, D3D11, D3D12) populates these.
Backend :: struct {
	capabilities:                  Capabilities,

	// --- Lifecycle ---
	shutdown:                      proc(),
	wait_idle:                     proc(),

	// --- Frame ---
	begin_frame:                   proc(clear_color: [4]f32) -> (Frame_Context, bool),
	end_frame:                     proc(ctx: Frame_Context) -> bool, // returns true if swapchain recreated
	on_resize:                     proc(width, height: u32),
	get_extent:                    proc() -> Extent,
	current_frame_index:           proc() -> u32,

	// --- Render pass ---
	begin_render_pass:             proc(ctx: Frame_Context, desc: Render_Pass_Begin_Desc),
	begin_default_pass:            proc(ctx: Frame_Context, clear_color: [4]f32),
	end_render_pass:               proc(ctx: Frame_Context),
	set_viewport:                  proc(ctx: Frame_Context, x, y, w, h: f32),
	set_scissor:                   proc(ctx: Frame_Context, x, y: i32, w, h: u32),
	set_depth_bias:                proc(ctx: Frame_Context, constant, slope: f32),

	// --- Pipeline ---
	create_graphics_pipeline:      proc(desc: Pipeline_Desc) -> (Pipeline_Handle, bool),
	destroy_graphics_pipeline:     proc(handle: Pipeline_Handle),
	bind_graphics_pipeline:        proc(ctx: Frame_Context, handle: Pipeline_Handle),
	push_constants:                proc(
		ctx: Frame_Context,
		pipeline: Pipeline_Handle,
		stages: Shader_Stage_Flags,
		offset, size: u32,
		data: rawptr,
	),

	// --- Buffers ---
	create_buffer:                 proc(desc: Buffer_Desc) -> (Buffer_Handle, bool),
	create_buffer_staged:          proc(
		data: rawptr,
		size: int,
		usage: Buffer_Usage_Flags,
	) -> (
		Buffer_Handle,
		bool,
	),
	destroy_buffer:                proc(handle: Buffer_Handle),
	map_buffer:                    proc(handle: Buffer_Handle) -> rawptr,
	unmap_buffer:                  proc(handle: Buffer_Handle),
	get_buffer_mapped:             proc(handle: Buffer_Handle) -> rawptr,
	bind_vertex_buffer:            proc(ctx: Frame_Context, handle: Buffer_Handle),
	bind_vertex_buffer_slot:       proc(
		ctx: Frame_Context,
		slot: u32,
		handle: Buffer_Handle,
		offset: u64,
		stride: u32,
	),
	bind_index_buffer:             proc(ctx: Frame_Context, handle: Buffer_Handle),

	// --- Textures ---
	create_texture:                proc(
		desc: Texture_Desc,
		pixels: rawptr,
	) -> (
		Texture_Handle,
		bool,
	),
	destroy_texture:               proc(handle: Texture_Handle),
	read_texture_rgba8:            proc(desc: Readback_Texture_Desc, out: []u8) -> bool,

	// --- Samplers ---
	create_sampler:                proc(desc: Sampler_Desc) -> (Sampler_Handle, bool),
	destroy_sampler:               proc(handle: Sampler_Handle),

	// --- Images (for depth/shadow targets) ---
	create_image:                  proc(desc: Texture_Desc) -> (Texture_Handle, bool),
	create_image_view:             proc(
		texture: Texture_Handle,
		format: Format,
		aspect: Image_Aspect_Flags,
	) -> bool,
	destroy_image:                 proc(handle: Texture_Handle),

	// --- Descriptors ---
	create_descriptor_set_layout:  proc(
		bindings: []Descriptor_Set_Layout_Binding,
	) -> (
		Descriptor_Handle,
		bool,
	),
	destroy_descriptor_set_layout: proc(handle: Descriptor_Handle),
	create_descriptor_pool:        proc(
		max_sets: u32,
		types: []Descriptor_Type,
		counts: []u32,
	) -> (
		Descriptor_Handle,
		bool,
	),
	destroy_descriptor_pool:       proc(handle: Descriptor_Handle),
	allocate_descriptor_set:       proc(
		pool: Descriptor_Handle,
		layout: Descriptor_Handle,
	) -> (
		Descriptor_Handle,
		bool,
	),
	bind_descriptor_set:           proc(
		ctx: Frame_Context,
		pipeline: Pipeline_Handle,
		set: Descriptor_Handle,
		index: u32,
	),
	update_descriptor_image:       proc(
		set: Descriptor_Handle,
		binding: u32,
		texture: Texture_Handle,
		sampler: Sampler_Handle,
		layout: Image_Layout,
	),
	update_descriptor_buffer:      proc(
		set: Descriptor_Handle,
		binding: u32,
		buffer: Buffer_Handle,
		size: u64,
	),

	// --- Render pass objects ---
	create_render_pass:            proc(desc: Render_Pass_Desc) -> (Render_Pass_Handle, bool),
	destroy_render_pass:           proc(handle: Render_Pass_Handle),
	create_framebuffer:            proc(desc: Framebuffer_Desc) -> (Framebuffer_Handle, bool),
	destroy_framebuffer:           proc(handle: Framebuffer_Handle),

	// --- Shaders ---
	create_shader_module:          proc(desc: Shader_Module_Desc) -> (Shader_Handle, bool),
	destroy_shader:                proc(handle: Shader_Handle),

	// --- Draw ---
	draw:                          proc(
		ctx: Frame_Context,
		vertex_count, instance_count: u32,
		first_vertex: u32,
		first_instance: u32,
	),
	draw_indexed:                  proc(
		ctx: Frame_Context,
		index_count, instance_count: u32,
		first_index: u32,
		vertex_offset: i32,
		first_instance: u32,
	),
	draw_indirect:                 proc(
		ctx: Frame_Context,
		argument_buffer: Buffer_Handle,
		argument_offset: u64,
		draw_count: u32,
		stride: u32,
	),
	draw_indexed_indirect:         proc(
		ctx: Frame_Context,
		argument_buffer: Buffer_Handle,
		argument_offset: u64,
		draw_count: u32,
		stride: u32,
	),

	// --- Compute ---
	create_compute_pipeline:       proc(
		shader: Shader_Handle,
		num_buffers: u32,
		push_constant_size: u32,
	) -> (
		Pipeline_Handle,
		bool,
	),
	destroy_compute_pipeline:      proc(handle: Pipeline_Handle),
	bind_compute_pipeline:         proc(ctx: Frame_Context, handle: Pipeline_Handle),
	dispatch_compute:              proc(ctx: Frame_Context, groups_x, groups_y, groups_z: u32),
	compute_barrier:               proc(ctx: Frame_Context),

	// --- Sync ---
	get_default_render_pass:       proc() -> Render_Pass_Handle,
	get_depth_format:              proc() -> Format,
}

backend_vtable_is_complete :: proc(b: ^Backend) -> bool {
	if b == nil do return false
	return(
		b.shutdown != nil &&
		b.wait_idle != nil &&
		b.begin_frame != nil &&
		b.end_frame != nil &&
		b.on_resize != nil &&
		b.get_extent != nil &&
		b.current_frame_index != nil &&
		b.begin_render_pass != nil &&
		b.begin_default_pass != nil &&
		b.end_render_pass != nil &&
		b.set_viewport != nil &&
		b.set_scissor != nil &&
		b.set_depth_bias != nil &&
		b.create_graphics_pipeline != nil &&
		b.destroy_graphics_pipeline != nil &&
		b.bind_graphics_pipeline != nil &&
		b.push_constants != nil &&
		b.create_buffer != nil &&
		b.create_buffer_staged != nil &&
		b.destroy_buffer != nil &&
		b.map_buffer != nil &&
		b.unmap_buffer != nil &&
		b.get_buffer_mapped != nil &&
		b.bind_vertex_buffer != nil &&
		b.bind_vertex_buffer_slot != nil &&
		b.bind_index_buffer != nil &&
		b.create_texture != nil &&
		b.destroy_texture != nil &&
		b.create_sampler != nil &&
		b.destroy_sampler != nil &&
		b.create_image != nil &&
		b.create_image_view != nil &&
		b.destroy_image != nil &&
		b.create_descriptor_set_layout != nil &&
		b.destroy_descriptor_set_layout != nil &&
		b.create_descriptor_pool != nil &&
		b.destroy_descriptor_pool != nil &&
		b.allocate_descriptor_set != nil &&
		b.bind_descriptor_set != nil &&
		b.update_descriptor_image != nil &&
		b.update_descriptor_buffer != nil &&
		b.create_render_pass != nil &&
		b.destroy_render_pass != nil &&
		b.create_framebuffer != nil &&
		b.destroy_framebuffer != nil &&
		b.create_shader_module != nil &&
		b.destroy_shader != nil &&
		b.draw != nil &&
		b.draw_indexed != nil &&
		b.draw_indirect != nil &&
		b.draw_indexed_indirect != nil &&
		b.create_compute_pipeline != nil &&
		b.destroy_compute_pipeline != nil &&
		b.bind_compute_pipeline != nil &&
		b.dispatch_compute != nil &&
		b.compute_barrier != nil &&
		b.get_default_render_pass != nil &&
		b.get_depth_format != nil \
	)
}

// Initialized flag
backend_initialized: bool