package vk_backend

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

// --- Lifecycle ---

shutdown_vk :: proc() {
	if g_vk == nil {return}

	vk.DeviceWaitIdle(g_vk.device.device)

	// Destroy all active samplers
	for &entry in g_vk.samplers {
		if entry.active {
			vk.DestroySampler(g_vk.device.device, entry.handle, nil)
			entry.active = false
		}
	}

	// Destroy all active framebuffers
	for &entry in g_vk.framebuffers {
		if entry.active {
			vk.DestroyFramebuffer(g_vk.device.device, entry.handle, nil)
			entry.active = false
		}
	}

	// Destroy all active render passes (except swapchain's)
	for &entry in g_vk.render_passes {
		if entry.active && entry.handle != g_vk.swapchain.render_pass {
			vk.DestroyRenderPass(g_vk.device.device, entry.handle, nil)
		}
		entry.active = false
	}

	// Destroy all active descriptors
	for &entry in g_vk.descriptors {
		if entry.active {
			switch entry.kind {
			case .Set_Layout:
				vk.DestroyDescriptorSetLayout(g_vk.device.device, entry.set_layout, nil)
			case .Pool:
				vk.DestroyDescriptorPool(g_vk.device.device, entry.pool, nil)
			case .Set:
			// Sets freed when pool is destroyed
			}
			entry.active = false
		}
	}

	// Destroy all active shaders
	for &entry in g_vk.shaders {
		if entry.active {
			vk.DestroyShaderModule(g_vk.device.device, entry.module.handle, nil)
			entry.active = false
		}
	}

	// Destroy all active pipelines
	for &entry in g_vk.pipelines {
		if entry.active {
			if entry.is_compute {
				if entry.compute.pipeline != 0 {
					vk.DestroyPipeline(g_vk.device.device, entry.compute.pipeline, nil)
				}
				if entry.compute.layout != 0 {
					vk.DestroyPipelineLayout(g_vk.device.device, entry.compute.layout, nil)
				}
				if entry.compute.desc_pool != 0 {
					vk.DestroyDescriptorPool(g_vk.device.device, entry.compute.desc_pool, nil)
				}
				if entry.compute.desc_set_layout != 0 {
					vk.DestroyDescriptorSetLayout(
						g_vk.device.device,
						entry.compute.desc_set_layout,
						nil,
					)
				}
			} else {
				if entry.graphics.pipeline != 0 {
					vk.DestroyPipeline(g_vk.device.device, entry.graphics.pipeline, nil)
				}
				if entry.graphics.layout != 0 {
					vk.DestroyPipelineLayout(g_vk.device.device, entry.graphics.layout, nil)
				}
			}
			entry.active = false
		}
	}

	// Destroy swapchain framebuffers before textures -- they reference the depth view
	for fb in g_vk.swapchain.framebuffers {
		if fb != 0 {
			vk.DestroyFramebuffer(g_vk.device.device, fb, nil)
		}
	}
	delete(g_vk.swapchain.framebuffers)
	g_vk.swapchain.framebuffers = nil

	// Destroy all active textures
	for &entry in g_vk.textures {
		if entry.active {
			gpu.destroy_image(&g_vk.device, &entry.image)
			entry.active = false
		}
	}

	// Clear swapchain depth_view since the texture pool already destroyed it
	g_vk.swapchain.depth_view = 0

	// Destroy all active buffers
	for &entry in g_vk.buffers {
		if entry.active {
			gpu.destroy_buffer(&g_vk.device, &entry.buffer)
			entry.active = false
		}
	}

	// Destroy core Vulkan objects
	if g_vk.cmd_pool != 0 {
		vk.DestroyCommandPool(g_vk.device.device, g_vk.cmd_pool, nil)
	}
	gpu.destroy_sync_objects(&g_vk.device, &g_vk.frame_sync)
	gpu.destroy_swapchain(&g_vk.swapchain, &g_vk.device)
	if g_vk.surface != 0 {
		vk.DestroySurfaceKHR(g_vk.instance.instance, g_vk.surface, nil)
	}
	gpu.destroy_device(&g_vk.device)
	gpu.destroy_instance(&g_vk.instance)

	free(g_vk)
	g_vk = nil
	bk.backend_initialized = false
}

wait_idle_vk :: proc() {
	if g_vk != nil {
		vk.DeviceWaitIdle(g_vk.device.device)
	}
}

// --- Frame lifecycle ---

begin_frame_vk :: proc(clear_color: [4]f32) -> (bk.Frame_Context, bool) {
	if g_vk == nil {return {}, false}

	frame := g_vk.current_frame
	sync := &g_vk.frame_sync[frame]
	cmd := g_vk.cmd_buffers[frame]

	// Wait for previous frame's fence
	vk.WaitForFences(g_vk.device.device, 1, &sync.in_flight, true, max(u64))

	if g_vk.headless {
		vk.ResetFences(g_vk.device.device, 1, &sync.in_flight)
		vk.ResetCommandBuffer(cmd, {})
		begin_info := vk.CommandBufferBeginInfo {
			sType = .COMMAND_BUFFER_BEGIN_INFO,
		}
		result := vk.BeginCommandBuffer(cmd, &begin_info)
		if result != .SUCCESS {
			log.errorf("gpu/vk: headless vkBeginCommandBuffer failed: %v", result)
			return {}, false
		}
		g_vk.frame_active = true
		return bk.Frame_Context{frame_index = frame}, true
	}

	// Acquire next swapchain image
	result := vk.AcquireNextImageKHR(
		g_vk.device.device,
		g_vk.swapchain.handle,
		max(u64),
		sync.image_available,
		0,
		&g_vk.image_index,
	)

	if result == .ERROR_OUT_OF_DATE_KHR {
		// Swapchain is stale — recreate immediately so the next frame can acquire.
		extent := g_vk.swapchain.extent
		gpu.recreate_swapchain(
			&g_vk.swapchain,
			&g_vk.device,
			g_vk.surface,
			extent.width,
			extent.height,
		)
		recreate_swapchain_depth_and_framebuffers()
		if g_vk.default_render_pass != bk.NULL_RENDER_PASS {
			g_vk.render_passes[g_vk.default_render_pass].handle = g_vk.swapchain.render_pass
		}
		return {}, false
	}

	// Reset fence after successful acquire
	vk.ResetFences(g_vk.device.device, 1, &sync.in_flight)

	// Reset and begin command buffer
	vk.ResetCommandBuffer(cmd, {})
	begin_info := vk.CommandBufferBeginInfo {
		sType = .COMMAND_BUFFER_BEGIN_INFO,
	}
	vk.BeginCommandBuffer(cmd, &begin_info)

	g_vk.frame_active = true

	return bk.Frame_Context{frame_index = frame}, true
}

end_frame_vk :: proc(ctx: bk.Frame_Context) -> bool {
	if g_vk == nil || !g_vk.frame_active {return false}

	frame := ctx.frame_index
	sync := &g_vk.frame_sync[frame]
	cmd := g_vk.cmd_buffers[frame]

	// End command buffer
	result := vk.EndCommandBuffer(cmd)
	if result != .SUCCESS {
		log.errorf("gpu/vk: vkEndCommandBuffer failed: %v", result)
		g_vk.frame_active = false
		return false
	}

	if g_vk.headless {
		submit_info := vk.SubmitInfo {
			sType              = .SUBMIT_INFO,
			commandBufferCount = 1,
			pCommandBuffers    = &cmd,
		}
		result = vk.QueueSubmit(g_vk.device.graphics_queue, 1, &submit_info, sync.in_flight)
		if result != .SUCCESS {
			log.errorf("gpu/vk: headless vkQueueSubmit failed: %v", result)
		}
		g_vk.frame_active = false
		g_vk.current_frame = (frame + 1) % bk.MAX_FRAMES_IN_FLIGHT
		return result == .SUCCESS
	}

	// Submit
	wait_stage := vk.PipelineStageFlags{.COLOR_ATTACHMENT_OUTPUT}
	submit_info := vk.SubmitInfo {
		sType                = .SUBMIT_INFO,
		waitSemaphoreCount   = 1,
		pWaitSemaphores      = &sync.image_available,
		pWaitDstStageMask    = &wait_stage,
		commandBufferCount   = 1,
		pCommandBuffers      = &cmd,
		signalSemaphoreCount = 1,
		pSignalSemaphores    = &g_vk.swapchain.render_finished_semaphores[g_vk.image_index],
	}

	result = vk.QueueSubmit(g_vk.device.graphics_queue, 1, &submit_info, sync.in_flight)
	if result != .SUCCESS {
		log.errorf("gpu/vk: vkQueueSubmit failed: %v", result)
	}

	// Present
	sc_handle := g_vk.swapchain.handle
	present_info := vk.PresentInfoKHR {
		sType              = .PRESENT_INFO_KHR,
		waitSemaphoreCount = 1,
		pWaitSemaphores    = &g_vk.swapchain.render_finished_semaphores[g_vk.image_index],
		swapchainCount     = 1,
		pSwapchains        = &sc_handle,
		pImageIndices      = &g_vk.image_index,
	}

	result = vk.QueuePresentKHR(g_vk.device.present_queue, &present_info)

	g_vk.frame_active = false
	g_vk.current_frame = (frame + 1) % bk.MAX_FRAMES_IN_FLIGHT

	if result == .ERROR_OUT_OF_DATE_KHR || result == .SUBOPTIMAL_KHR {
		// Wait for the just-submitted GPU work to complete before
		// destroying swapchain resources. This ensures semaphores are
		// fully consumed and Wayland explicit sync state is clean.
		vk.WaitForFences(g_vk.device.device, 1, &sync.in_flight, true, max(u64))

		extent := g_vk.swapchain.extent
		gpu.recreate_swapchain(
			&g_vk.swapchain,
			&g_vk.device,
			g_vk.surface,
			extent.width,
			extent.height,
		)
		recreate_swapchain_depth_and_framebuffers()
		// Update default render pass handle
		if g_vk.default_render_pass != bk.NULL_RENDER_PASS {
			g_vk.render_passes[g_vk.default_render_pass].handle = g_vk.swapchain.render_pass
		}
		return true
	}

	return false
}

on_resize_vk :: proc(width, height: u32) {
	if g_vk == nil {return}
	if width == 0 || height == 0 {return}
	if g_vk.headless {
		g_vk.headless_extent = {width, height}
		return
	}

	vk.DeviceWaitIdle(g_vk.device.device)
	gpu.recreate_swapchain(&g_vk.swapchain, &g_vk.device, g_vk.surface, width, height)
	recreate_swapchain_depth_and_framebuffers()

	// Update default render pass handle
	if g_vk.default_render_pass != bk.NULL_RENDER_PASS {
		g_vk.render_passes[g_vk.default_render_pass].handle = g_vk.swapchain.render_pass
	}
}

get_extent_vk :: proc() -> bk.Extent {
	if g_vk == nil {return {}}
	if g_vk.headless do return g_vk.headless_extent
	return bk.Extent{width = g_vk.swapchain.extent.width, height = g_vk.swapchain.extent.height}
}

// --- Sync queries ---

get_default_render_pass_vk :: proc() -> bk.Render_Pass_Handle {
	if g_vk == nil {return bk.NULL_RENDER_PASS}
	return g_vk.default_render_pass
}

get_depth_format_vk :: proc() -> bk.Format {
	if g_vk == nil {return .Undefined}
	if g_vk.headless {return .Undefined}
	return from_vk_format(g_vk.swapchain.depth_format)
}

// Recreate depth buffer and framebuffers after swapchain recreation.
@(private)
recreate_swapchain_depth_and_framebuffers :: proc() {
	sc := &g_vk.swapchain

	// Destroy old depth texture pool entry (image + view + memory)
	if sc.depth_view != 0 {
		for &entry in g_vk.textures {
			if entry.active && entry.image.view == sc.depth_view {
				gpu.destroy_image(&g_vk.device, &entry.image)
				entry.active = false
				break
			}
		}
		sc.depth_view = 0
	}

	// Recreate depth buffer if we have a depth format
	if sc.depth_format != .UNDEFINED {
		depth_img, depth_ok := gpu.create_image(
			&g_vk.device,
			sc.extent.width,
			sc.extent.height,
			sc.depth_format,
			.OPTIMAL,
			{.DEPTH_STENCIL_ATTACHMENT},
			{.DEVICE_LOCAL},
		)
		if depth_ok {
			depth_view, dv_ok := gpu.create_image_view(
				&g_vk.device,
				depth_img.image,
				sc.depth_format,
				{.DEPTH},
			)
			if dv_ok {
				depth_img.view = depth_view
				sc.depth_view = depth_view
				// Store in texture pool
				dh, dh_ok := alloc_texture_handle()
				if dh_ok {
					g_vk.textures[dh].image = depth_img
					g_vk.textures[dh].active = true
				}
			} else {
				gpu.destroy_image(&g_vk.device, &depth_img)
			}
		}
	}

	gpu.create_framebuffers(sc, &g_vk.device)
}