package vk_backend

import bk ".."
import gpu "../../core"
import "../../pipeline"
import "core:log"
import "core:os"
import vk "vendor:vulkan"

// --- Global Vulkan state ---

@(private)
g_vk: ^Vulkan_State

Vulkan_State :: struct {
	// Core
	instance:            gpu.Gpu_Instance,
	surface:             vk.SurfaceKHR,
	device:              gpu.Gpu_Device,
	swapchain:           gpu.Swapchain,
	cmd_pool:            vk.CommandPool,
	cmd_buffers:         [bk.MAX_FRAMES_IN_FLIGHT]vk.CommandBuffer,
	frame_sync:          [bk.MAX_FRAMES_IN_FLIGHT]gpu.Frame_Sync,
	current_frame:       u32,
	headless:            bool,
	headless_extent:     bk.Extent,

	// Handle pools
	buffers:             [MAX_BUFFERS]Vk_Buffer_Entry,
	textures:            [MAX_TEXTURES]Vk_Texture_Entry,
	pipelines:           [MAX_PIPELINES]Vk_Pipeline_Entry,
	shaders:             [MAX_SHADERS]Vk_Shader_Entry,
	descriptors:         [MAX_DESCRIPTORS]Vk_Descriptor_Entry,
	render_passes:       [MAX_RENDER_PASSES]Vk_Render_Pass_Entry,
	framebuffers:        [MAX_FRAMEBUFFERS]Vk_Framebuffer_Entry,
	samplers:            [MAX_SAMPLERS]Vk_Sampler_Entry,

	// Current frame state
	image_index:         u32,
	frame_active:        bool,

	// Default render pass handle (swapchain render pass)
	default_render_pass: bk.Render_Pass_Handle,
}

// Pool sizes
MAX_BUFFERS :: 1024
MAX_TEXTURES :: 512
MAX_PIPELINES :: 128
MAX_SHADERS :: 256
MAX_DESCRIPTORS :: 512
MAX_RENDER_PASSES :: 32
MAX_FRAMEBUFFERS :: 64
MAX_SAMPLERS :: 64

VULKAN_IMPLEMENTED_PUSH_CONSTANT_MAX_SIZE :: 256

// Pool entry types
Vk_Buffer_Entry :: struct {
	buffer: gpu.Gpu_Buffer,
	memory: bk.Memory_Property_Flags,
	active: bool,
}

Vk_Texture_Entry :: struct {
	image:  gpu.Gpu_Image,
	active: bool,
}

Vk_Pipeline_Entry :: struct {
	// Can be graphics or compute
	graphics:   pipeline.Graphics_Pipeline,
	compute:    pipeline.Compute_Pipeline,
	is_compute: bool,
	active:     bool,
}

Vk_Shader_Entry :: struct {
	module: pipeline.Shader_Module,
	active: bool,
}

Vk_Descriptor_Entry :: struct {
	// Can be set layout, pool, or set
	set_layout:      vk.DescriptorSetLayout,
	pool:            vk.DescriptorPool,
	set:             vk.DescriptorSet,
	layout_bindings: [16]bk.Descriptor_Set_Layout_Binding,
	layout_count:    u32,
	bindings:        [16]Vk_Descriptor_Binding,
	binding_count:   u32,
	kind:            Descriptor_Kind,
	active:          bool,
}

Vk_Descriptor_Binding :: struct {
	binding: u32,
	type:    bk.Descriptor_Type,
}

Descriptor_Kind :: enum {
	Set_Layout,
	Pool,
	Set,
}

Vk_Render_Pass_Entry :: struct {
	handle: vk.RenderPass,
	desc:   bk.Render_Pass_Desc,
	active: bool,
}

Vk_Framebuffer_Entry :: struct {
	handle: vk.Framebuffer,
	active: bool,
}

Vk_Sampler_Entry :: struct {
	handle: vk.Sampler,
	active: bool,
}

// --- Init ---

init_vulkan_backend :: proc(
	surface_desc: bk.Surface_Desc,
	width, height: u32,
	title: cstring,
) -> (
	backend: bk.Backend,
	ok: bool,
) {
	state := new(Vulkan_State)
	if state == nil {
		log.error("gpu/vk: failed to allocate vulkan state")
		return {}, false
	}
	g_vk = state

	enable_validation := os.get_env("GPU_VK_VALIDATION", context.temp_allocator) != "0"
	inst, inst_ok := gpu.create_instance(title, enable_validation)
	if !inst_ok {
		log.error("gpu/vk: failed to create Vulkan instance")
		free(state)
		g_vk = nil
		return {}, false
	}
	state.instance = inst
	state.headless = surface_desc.kind == .None
	state.headless_extent = {width, height}

	if state.headless {
		dev, dev_ok := gpu.create_device_headless(inst.instance)
		if !dev_ok {
			log.error("gpu/vk: failed to create headless Vulkan device")
			gpu.destroy_instance(&state.instance)
			free(state)
			g_vk = nil
			return {}, false
		}
		state.device = dev
		if !init_frame_resources_vk(state) {
			return {}, false
		}
		backend = make_backend_vk(state)
		bk.backend_initialized = true
		return backend, true
	}

	// Create surface
	surface, surface_ok := gpu.create_surface(surface_desc, inst.instance)
	if !surface_ok {
		log.error("gpu/vk: failed to create Vulkan surface")
		gpu.destroy_instance(&state.instance)
		free(state)
		g_vk = nil
		return {}, false
	}
	state.surface = surface

	// Create device
	dev, dev_ok := gpu.create_device(inst.instance, surface)
	if !dev_ok {
		log.error("gpu/vk: failed to create Vulkan device")
		vk.DestroySurfaceKHR(inst.instance, surface, nil)
		gpu.destroy_instance(&state.instance)
		free(state)
		g_vk = nil
		return {}, false
	}
	state.device = dev

	// Find depth format
	depth_fmt := find_depth_format_vk(&state.device)

	// Create swapchain
	sc, sc_ok := gpu.create_swapchain(&state.device, surface, width, height, depth_fmt)
	if !sc_ok {
		log.error("gpu/vk: failed to create swapchain")
		gpu.destroy_device(&state.device)
		vk.DestroySurfaceKHR(inst.instance, surface, nil)
		gpu.destroy_instance(&state.instance)
		free(state)
		g_vk = nil
		return {}, false
	}
	state.swapchain = sc

	// Create depth buffer and framebuffers
	if depth_fmt != .UNDEFINED {
		depth_img, depth_ok := gpu.create_image(
			&state.device,
			sc.extent.width,
			sc.extent.height,
			depth_fmt,
			.OPTIMAL,
			{.DEPTH_STENCIL_ATTACHMENT},
			{.DEVICE_LOCAL},
		)
		if !depth_ok {
			log.error("gpu/vk: failed to create depth image")
			gpu.destroy_swapchain(&state.swapchain, &state.device)
			gpu.destroy_device(&state.device)
			vk.DestroySurfaceKHR(inst.instance, surface, nil)
			gpu.destroy_instance(&state.instance)
			free(state)
			g_vk = nil
			return {}, false
		}

		depth_view, dv_ok := gpu.create_image_view(
			&state.device,
			depth_img.image,
			depth_fmt,
			{.DEPTH},
		)
		if !dv_ok {
			log.error("gpu/vk: failed to create depth image view")
			gpu.destroy_image(&state.device, &depth_img)
			gpu.destroy_swapchain(&state.swapchain, &state.device)
			gpu.destroy_device(&state.device)
			vk.DestroySurfaceKHR(inst.instance, surface, nil)
			gpu.destroy_instance(&state.instance)
			free(state)
			g_vk = nil
			return {}, false
		}

		depth_img.view = depth_view
		state.swapchain.depth_view = depth_view

		// Store depth image in texture pool so it gets cleaned up
		dh, dh_ok := alloc_texture_handle()
		if dh_ok {
			state.textures[dh].image = depth_img
			state.textures[dh].active = true
		}
	}

	if !gpu.create_framebuffers(&state.swapchain, &state.device) {
		log.error("gpu/vk: failed to create framebuffers")
		gpu.destroy_swapchain(&state.swapchain, &state.device)
		gpu.destroy_device(&state.device)
		vk.DestroySurfaceKHR(inst.instance, surface, nil)
		gpu.destroy_instance(&state.instance)
		free(state)
		g_vk = nil
		return {}, false
	}

	if !init_frame_resources_vk(state) {
		return {}, false
	}

	// Register default render pass (swapchain render pass) in pool
	rp_handle, rp_ok := alloc_render_pass_handle()
	if !rp_ok {
		log.error("gpu/vk: failed to allocate default render pass handle")
		shutdown_vk()
		return {}, false
	}
	state.render_passes[rp_handle].handle = state.swapchain.render_pass
	state.render_passes[rp_handle].active = true
	state.default_render_pass = rp_handle

	// Populate backend vtable
	backend = make_backend_vk(state)

	bk.backend_initialized = true
	return backend, true
}

@(private)
init_frame_resources_vk :: proc(state: ^Vulkan_State) -> bool {
	pool, pool_ok := gpu.create_command_pool(&state.device, state.device.queue_families.graphics)
	if !pool_ok {
		log.error("gpu/vk: failed to create command pool")
		shutdown_vk()
		return false
	}
	state.cmd_pool = pool

	bufs, bufs_ok := gpu.allocate_command_buffers(&state.device, pool, bk.MAX_FRAMES_IN_FLIGHT)
	if !bufs_ok {
		log.error("gpu/vk: failed to allocate command buffers")
		shutdown_vk()
		return false
	}
	for i in 0 ..< bk.MAX_FRAMES_IN_FLIGHT {
		state.cmd_buffers[i] = bufs[i]
	}

	sync, sync_ok := gpu.create_sync_objects(&state.device)
	if !sync_ok {
		log.error("gpu/vk: failed to create sync objects")
		shutdown_vk()
		return false
	}
	state.frame_sync = sync
	return true
}

@(private)
make_backend_vk :: proc(state: ^Vulkan_State) -> bk.Backend {
	return bk.Backend {
		capabilities                  = bk.implemented_base_capabilities(
			bk.MAX_COLOR_TARGETS,
			min(
				state.device.properties.limits.maxPushConstantsSize,
				u32(VULKAN_IMPLEMENTED_PUSH_CONSTANT_MAX_SIZE),
			),
		),
		shutdown                      = shutdown_vk,
		wait_idle                     = wait_idle_vk,
		begin_frame                   = begin_frame_vk,
		end_frame                     = end_frame_vk,
		on_resize                     = on_resize_vk,
		get_extent                    = get_extent_vk,
		current_frame_index           = get_current_frame,
		begin_render_pass             = begin_render_pass_vk,
		begin_default_pass            = begin_default_pass_vk,
		end_render_pass               = end_render_pass_vk,
		set_viewport                  = set_viewport_vk,
		set_scissor                   = set_scissor_vk,
		set_depth_bias                = set_depth_bias_vk,
		create_graphics_pipeline      = create_graphics_pipeline_vk,
		destroy_graphics_pipeline     = destroy_graphics_pipeline_vk,
		bind_graphics_pipeline        = bind_graphics_pipeline_vk,
		push_constants                = push_constants_vk,
		create_buffer                 = create_buffer_vk,
		create_buffer_staged          = create_buffer_staged_vk,
		destroy_buffer                = destroy_buffer_vk,
		map_buffer                    = map_buffer_vk,
		unmap_buffer                  = unmap_buffer_vk,
		get_buffer_mapped             = get_buffer_mapped_vk,
		bind_vertex_buffer            = bind_vertex_buffer_vk,
		bind_vertex_buffer_slot       = bind_vertex_buffer_slot_vk,
		bind_index_buffer             = bind_index_buffer_vk,
		create_texture                = create_texture_vk,
		destroy_texture               = destroy_texture_vk,
		read_texture_rgba8            = read_texture_rgba8_vk,
		create_sampler                = create_sampler_vk,
		destroy_sampler               = destroy_sampler_vk,
		create_image                  = create_image_vk,
		create_image_view             = create_image_view_vk,
		destroy_image                 = destroy_image_vk,
		create_descriptor_set_layout  = create_descriptor_set_layout_vk,
		destroy_descriptor_set_layout = destroy_descriptor_set_layout_vk,
		create_descriptor_pool        = create_descriptor_pool_vk,
		destroy_descriptor_pool       = destroy_descriptor_pool_vk,
		allocate_descriptor_set       = allocate_descriptor_set_vk,
		bind_descriptor_set           = bind_descriptor_set_vk,
		update_descriptor_image       = update_descriptor_image_vk,
		update_descriptor_buffer      = update_descriptor_buffer_vk,
		create_render_pass            = create_render_pass_vk,
		destroy_render_pass           = destroy_render_pass_vk,
		create_framebuffer            = create_framebuffer_vk,
		destroy_framebuffer           = destroy_framebuffer_vk,
		create_shader_module          = create_shader_module_vk,
		destroy_shader                = destroy_shader_vk,
		draw                          = draw_vk,
		draw_indexed                  = draw_indexed_vk,
		draw_indirect                 = draw_indirect_vk,
		draw_indexed_indirect         = draw_indexed_indirect_vk,
		create_compute_pipeline       = create_compute_pipeline_vk,
		destroy_compute_pipeline      = destroy_compute_pipeline_vk,
		bind_compute_pipeline         = bind_compute_pipeline_vk,
		dispatch_compute              = dispatch_compute_vk,
		compute_barrier               = compute_barrier_vk,
		get_default_render_pass       = get_default_render_pass_vk,
		get_depth_format              = get_depth_format_vk,
	}
}

// --- Helper: get Vulkan command buffer from frame context ---

@(private)
get_cmd :: proc(ctx: bk.Frame_Context) -> vk.CommandBuffer {
	return g_vk.cmd_buffers[ctx.frame_index]
}

// --- Handle allocation helpers ---

@(private)
alloc_buffer_handle :: proc() -> (bk.Buffer_Handle, bool) {
	for i in 1 ..< u64(MAX_BUFFERS) {
		if !g_vk.buffers[i].active {
			return bk.Buffer_Handle(i), true
		}
	}
	return bk.NULL_BUFFER, false
}

@(private)
alloc_texture_handle :: proc() -> (bk.Texture_Handle, bool) {
	for i in 1 ..< u64(MAX_TEXTURES) {
		if !g_vk.textures[i].active {
			return bk.Texture_Handle(i), true
		}
	}
	return bk.NULL_TEXTURE, false
}

@(private)
alloc_pipeline_handle :: proc() -> (bk.Pipeline_Handle, bool) {
	for i in 1 ..< u64(MAX_PIPELINES) {
		if !g_vk.pipelines[i].active {
			return bk.Pipeline_Handle(i), true
		}
	}
	return bk.NULL_PIPELINE, false
}

@(private)
alloc_shader_handle :: proc() -> (bk.Shader_Handle, bool) {
	for i in 1 ..< u64(MAX_SHADERS) {
		if !g_vk.shaders[i].active {
			return bk.Shader_Handle(i), true
		}
	}
	return bk.NULL_SHADER, false
}

@(private)
alloc_descriptor_handle :: proc() -> (bk.Descriptor_Handle, bool) {
	for i in 1 ..< u64(MAX_DESCRIPTORS) {
		if !g_vk.descriptors[i].active {
			return bk.Descriptor_Handle(i), true
		}
	}
	return bk.NULL_DESCRIPTOR, false
}

@(private)
alloc_render_pass_handle :: proc() -> (bk.Render_Pass_Handle, bool) {
	for i in 1 ..< u64(MAX_RENDER_PASSES) {
		if !g_vk.render_passes[i].active {
			return bk.Render_Pass_Handle(i), true
		}
	}
	return bk.NULL_RENDER_PASS, false
}

@(private)
alloc_framebuffer_handle :: proc() -> (bk.Framebuffer_Handle, bool) {
	for i in 1 ..< u64(MAX_FRAMEBUFFERS) {
		if !g_vk.framebuffers[i].active {
			return bk.Framebuffer_Handle(i), true
		}
	}
	return bk.NULL_FRAMEBUFFER, false
}

@(private)
alloc_sampler_handle :: proc() -> (bk.Sampler_Handle, bool) {
	for i in 1 ..< u64(MAX_SAMPLERS) {
		if !g_vk.samplers[i].active {
			return bk.Sampler_Handle(i), true
		}
	}
	return bk.NULL_SAMPLER, false
}

// --- Depth format ---

@(private)
find_depth_format_vk :: proc(dev: ^gpu.Gpu_Device) -> vk.Format {
	candidates := [?]vk.Format{.D32_SFLOAT, .D32_SFLOAT_S8_UINT, .D24_UNORM_S8_UINT}
	for fmt in candidates {
		props: vk.FormatProperties
		vk.GetPhysicalDeviceFormatProperties(dev.physical_device, fmt, &props)
		if .DEPTH_STENCIL_ATTACHMENT in props.optimalTilingFeatures {
			return fmt
		}
	}
	return .D32_SFLOAT
}

// --- Expose internal state for legacy bridge (temporary, removed after full refactor) ---

get_device :: proc() -> ^gpu.Gpu_Device {
	if g_vk == nil {return nil}
	return &g_vk.device
}

get_swapchain :: proc() -> ^gpu.Swapchain {
	if g_vk == nil {return nil}
	return &g_vk.swapchain
}

get_cmd_pool :: proc() -> vk.CommandPool {
	if g_vk == nil {return 0}
	return g_vk.cmd_pool
}

get_surface :: proc() -> vk.SurfaceKHR {
	if g_vk == nil {return 0}
	return g_vk.surface
}

get_frame_sync :: proc() -> ^[bk.MAX_FRAMES_IN_FLIGHT]gpu.Frame_Sync {
	if g_vk == nil {return nil}
	return &g_vk.frame_sync
}

get_cmd_buffers :: proc() -> ^[bk.MAX_FRAMES_IN_FLIGHT]vk.CommandBuffer {
	if g_vk == nil {return nil}
	return &g_vk.cmd_buffers
}

get_current_frame :: proc() -> u32 {
	if g_vk == nil {return 0}
	return g_vk.current_frame
}

set_current_frame :: proc(frame: u32) {
	if g_vk != nil {
		g_vk.current_frame = frame
	}
}

get_image_index :: proc() -> u32 {
	if g_vk == nil {return 0}
	return g_vk.image_index
}

set_image_index :: proc(idx: u32) {
	if g_vk != nil {
		g_vk.image_index = idx
	}
}