package gpu_core

import "core:log"
import "core:math"
import vk "vendor:vulkan"

Swapchain :: struct {
	handle:      vk.SwapchainKHR,
	images:      []vk.Image,
	image_views: []vk.ImageView,
	format:      vk.SurfaceFormatKHR,
	extent:      vk.Extent2D,
	present_mode: vk.PresentModeKHR,

	// Render pass and framebuffers
	render_pass:  vk.RenderPass,
	framebuffers: []vk.Framebuffer,

	// Depth
	depth_format: vk.Format,
	depth_view:   vk.ImageView, // owned externally (by renderer)

	// Per-swapchain-image semaphores for presentation sync
	render_finished_semaphores: []vk.Semaphore,
}

create_swapchain :: proc(dev: ^Gpu_Device, surface: vk.SurfaceKHR, width, height: u32, depth_format: vk.Format = .UNDEFINED) -> (sc: Swapchain, ok: bool) {
	// Query surface capabilities
	caps: vk.SurfaceCapabilitiesKHR
	vk.GetPhysicalDeviceSurfaceCapabilitiesKHR(dev.physical_device, surface, &caps)

	// Choose format
	sc.format = choose_surface_format(dev.physical_device, surface)

	// Choose present mode
	sc.present_mode = choose_present_mode(dev.physical_device, surface)

	// Choose extent
	sc.extent = choose_extent(caps, width, height)

	// Image count (prefer triple buffering)
	image_count := caps.minImageCount + 1
	if caps.maxImageCount > 0 && image_count > caps.maxImageCount {
		image_count = caps.maxImageCount
	}

	create_info := vk.SwapchainCreateInfoKHR{
		sType            = .SWAPCHAIN_CREATE_INFO_KHR,
		surface          = surface,
		minImageCount    = image_count,
		imageFormat      = sc.format.format,
		imageColorSpace  = sc.format.colorSpace,
		imageExtent      = sc.extent,
		imageArrayLayers = 1,
		imageUsage       = {.COLOR_ATTACHMENT},
		preTransform     = caps.currentTransform,
		compositeAlpha   = {.OPAQUE},
		presentMode      = sc.present_mode,
		clipped          = true,
	}

	// Queue family sharing
	if dev.queue_families.graphics != dev.queue_families.present {
		indices := [2]u32{dev.queue_families.graphics, dev.queue_families.present}
		create_info.imageSharingMode = .CONCURRENT
		create_info.queueFamilyIndexCount = 2
		create_info.pQueueFamilyIndices = &indices[0]
	} else {
		create_info.imageSharingMode = .EXCLUSIVE
	}

	result := vk.CreateSwapchainKHR(dev.device, &create_info, nil, &sc.handle)
	if result != .SUCCESS {
		log.errorf("gpu/core: vkCreateSwapchainKHR failed: %v", result)
		return {}, false
	}

	// Get images
	vk.GetSwapchainImagesKHR(dev.device, sc.handle, &image_count, nil)
	sc.images = make([]vk.Image, image_count)
	vk.GetSwapchainImagesKHR(dev.device, sc.handle, &image_count, raw_data(sc.images))

	// Create image views
	sc.image_views = make([]vk.ImageView, image_count)
	for i in 0..<image_count {
		view_info := vk.ImageViewCreateInfo{
			sType    = .IMAGE_VIEW_CREATE_INFO,
			image    = sc.images[i],
			viewType = .D2,
			format   = sc.format.format,
			components = {
				r = .IDENTITY,
				g = .IDENTITY,
				b = .IDENTITY,
				a = .IDENTITY,
			},
			subresourceRange = {
				aspectMask     = {.COLOR},
				baseMipLevel   = 0,
				levelCount     = 1,
				baseArrayLayer = 0,
				layerCount     = 1,
			},
		}
		result = vk.CreateImageView(dev.device, &view_info, nil, &sc.image_views[i])
		if result != .SUCCESS {
			log.errorf("gpu/core: failed to create image view %d: %v", i, result)
			destroy_swapchain(&sc, dev)
			return {}, false
		}
	}

	// Create render pass
	sc.depth_format = depth_format
	rp, rp_ok := create_render_pass(dev, sc.format.format, depth_format)
	if !rp_ok {
		destroy_swapchain(&sc, dev)
		return {}, false
	}
	sc.render_pass = rp

	// Framebuffers are NOT created here -- caller must set sc.depth_view
	// (if using depth) and call create_framebuffers() after.

	// Create per-image presentation semaphores
	if !create_present_semaphores(&sc, dev) {
		destroy_swapchain(&sc, dev)
		return {}, false
	}

	return sc, true
}

recreate_swapchain :: proc(sc: ^Swapchain, dev: ^Gpu_Device, surface: vk.SurfaceKHR, width, height: u32) -> bool {
	if width == 0 || height == 0 {
		return true // minimized, skip
	}

	vk.DeviceWaitIdle(dev.device)

	// Save render pass (reuse it)
	old_rp := sc.render_pass

	// Cleanup old swapchain resources (but not render pass)
	cleanup_swapchain_resources(sc, dev)

	// Query new capabilities
	caps: vk.SurfaceCapabilitiesKHR
	vk.GetPhysicalDeviceSurfaceCapabilitiesKHR(dev.physical_device, surface, &caps)

	sc.extent = choose_extent(caps, width, height)

	image_count := caps.minImageCount + 1
	if caps.maxImageCount > 0 && image_count > caps.maxImageCount {
		image_count = caps.maxImageCount
	}

	old_handle := sc.handle

	create_info := vk.SwapchainCreateInfoKHR{
		sType            = .SWAPCHAIN_CREATE_INFO_KHR,
		surface          = surface,
		minImageCount    = image_count,
		imageFormat      = sc.format.format,
		imageColorSpace  = sc.format.colorSpace,
		imageExtent      = sc.extent,
		imageArrayLayers = 1,
		imageUsage       = {.COLOR_ATTACHMENT},
		preTransform     = caps.currentTransform,
		compositeAlpha   = {.OPAQUE},
		presentMode      = sc.present_mode,
		clipped          = true,
		oldSwapchain     = old_handle,
	}

	if dev.queue_families.graphics != dev.queue_families.present {
		indices := [2]u32{dev.queue_families.graphics, dev.queue_families.present}
		create_info.imageSharingMode = .CONCURRENT
		create_info.queueFamilyIndexCount = 2
		create_info.pQueueFamilyIndices = &indices[0]
	} else {
		create_info.imageSharingMode = .EXCLUSIVE
	}

	result := vk.CreateSwapchainKHR(dev.device, &create_info, nil, &sc.handle)

	if old_handle != 0 {
		vk.DestroySwapchainKHR(dev.device, old_handle, nil)
	}

	if result != .SUCCESS {
		log.errorf("gpu/core: swapchain recreation failed: %v", result)
		return false
	}

	// Get new images
	vk.GetSwapchainImagesKHR(dev.device, sc.handle, &image_count, nil)
	sc.images = make([]vk.Image, image_count)
	vk.GetSwapchainImagesKHR(dev.device, sc.handle, &image_count, raw_data(sc.images))

	// Create new image views
	sc.image_views = make([]vk.ImageView, image_count)
	for i in 0..<image_count {
		view_info := vk.ImageViewCreateInfo{
			sType    = .IMAGE_VIEW_CREATE_INFO,
			image    = sc.images[i],
			viewType = .D2,
			format   = sc.format.format,
			components = {r = .IDENTITY, g = .IDENTITY, b = .IDENTITY, a = .IDENTITY},
			subresourceRange = {
				aspectMask     = {.COLOR},
				baseMipLevel   = 0,
				levelCount     = 1,
				baseArrayLayer = 0,
				layerCount     = 1,
			},
		}
		result = vk.CreateImageView(dev.device, &view_info, nil, &sc.image_views[i])
		if result != .SUCCESS {
			return false
		}
	}

	sc.render_pass = old_rp

	// Framebuffers are NOT created here -- caller must recreate depth buffer,
	// set sc.depth_view, and call create_framebuffers() after.

	return create_present_semaphores(sc, dev)
}

destroy_swapchain :: proc(sc: ^Swapchain, dev: ^Gpu_Device) {
	cleanup_swapchain_resources(sc, dev)

	if sc.render_pass != 0 {
		vk.DestroyRenderPass(dev.device, sc.render_pass, nil)
		sc.render_pass = 0
	}

	if sc.handle != 0 {
		vk.DestroySwapchainKHR(dev.device, sc.handle, nil)
		sc.handle = 0
	}
}

@(private)
cleanup_swapchain_resources :: proc(sc: ^Swapchain, dev: ^Gpu_Device) {
	for fb in sc.framebuffers {
		if fb != 0 {
			vk.DestroyFramebuffer(dev.device, fb, nil)
		}
	}
	delete(sc.framebuffers)
	sc.framebuffers = nil

	for iv in sc.image_views {
		if iv != 0 {
			vk.DestroyImageView(dev.device, iv, nil)
		}
	}
	delete(sc.image_views)
	sc.image_views = nil

	delete(sc.images)
	sc.images = nil

	destroy_present_semaphores(sc, dev)
}

create_present_semaphores :: proc(sc: ^Swapchain, dev: ^Gpu_Device) -> bool {
	count := len(sc.images)
	sc.render_finished_semaphores = make([]vk.Semaphore, count)
	sem_info := vk.SemaphoreCreateInfo{
		sType = .SEMAPHORE_CREATE_INFO,
	}
	for i in 0..<count {
		result := vk.CreateSemaphore(dev.device, &sem_info, nil, &sc.render_finished_semaphores[i])
		if result != .SUCCESS {
			log.errorf("gpu/core: failed to create present semaphore %d: %v", i, result)
			destroy_present_semaphores(sc, dev)
			return false
		}
	}
	return true
}

destroy_present_semaphores :: proc(sc: ^Swapchain, dev: ^Gpu_Device) {
	for &s in sc.render_finished_semaphores {
		if s != 0 {
			vk.DestroySemaphore(dev.device, s, nil)
			s = 0
		}
	}
	delete(sc.render_finished_semaphores)
	sc.render_finished_semaphores = nil
}

@(private)
create_render_pass :: proc(dev: ^Gpu_Device, color_format: vk.Format, depth_format: vk.Format = .UNDEFINED) -> (vk.RenderPass, bool) {
	has_depth := depth_format != .UNDEFINED

	color_attachment := vk.AttachmentDescription{
		format         = color_format,
		samples        = {._1},
		loadOp         = .CLEAR,
		storeOp        = .STORE,
		stencilLoadOp  = .DONT_CARE,
		stencilStoreOp = .DONT_CARE,
		initialLayout  = .UNDEFINED,
		finalLayout    = .PRESENT_SRC_KHR,
	}

	depth_attachment := vk.AttachmentDescription{
		format         = depth_format,
		samples        = {._1},
		loadOp         = .CLEAR,
		storeOp        = .DONT_CARE,
		stencilLoadOp  = .DONT_CARE,
		stencilStoreOp = .DONT_CARE,
		initialLayout  = .UNDEFINED,
		finalLayout    = .DEPTH_STENCIL_ATTACHMENT_OPTIMAL,
	}

	color_ref := vk.AttachmentReference{
		attachment = 0,
		layout     = .COLOR_ATTACHMENT_OPTIMAL,
	}

	depth_ref := vk.AttachmentReference{
		attachment = 1,
		layout     = .DEPTH_STENCIL_ATTACHMENT_OPTIMAL,
	}

	subpass := vk.SubpassDescription{
		pipelineBindPoint       = .GRAPHICS,
		colorAttachmentCount    = 1,
		pColorAttachments       = &color_ref,
		pDepthStencilAttachment = &depth_ref if has_depth else nil,
	}

	dependency := vk.SubpassDependency{
		srcSubpass    = vk.SUBPASS_EXTERNAL,
		dstSubpass    = 0,
		srcStageMask  = {.COLOR_ATTACHMENT_OUTPUT, .EARLY_FRAGMENT_TESTS},
		srcAccessMask = {},
		dstStageMask  = {.COLOR_ATTACHMENT_OUTPUT, .EARLY_FRAGMENT_TESTS},
		dstAccessMask = {.COLOR_ATTACHMENT_WRITE, .DEPTH_STENCIL_ATTACHMENT_WRITE},
	}

	attachments: [2]vk.AttachmentDescription
	attachments[0] = color_attachment
	attachments[1] = depth_attachment
	attachment_count: u32 = 2 if has_depth else 1

	rp_info := vk.RenderPassCreateInfo{
		sType           = .RENDER_PASS_CREATE_INFO,
		attachmentCount = attachment_count,
		pAttachments    = &attachments[0],
		subpassCount    = 1,
		pSubpasses      = &subpass,
		dependencyCount = 1,
		pDependencies   = &dependency,
	}

	render_pass: vk.RenderPass
	result := vk.CreateRenderPass(dev.device, &rp_info, nil, &render_pass)
	if result != .SUCCESS {
		log.errorf("gpu/core: vkCreateRenderPass failed: %v", result)
		return 0, false
	}

	return render_pass, true
}

create_framebuffers :: proc(sc: ^Swapchain, dev: ^Gpu_Device) -> bool {
	sc.framebuffers = make([]vk.Framebuffer, len(sc.image_views))
	has_depth := sc.depth_view != 0

	for iv, i in sc.image_views {
		attachment_buf: [2]vk.ImageView
		attachment_buf[0] = iv
		attachment_buf[1] = sc.depth_view
		attachment_count: u32 = 2 if has_depth else 1

		fb_info := vk.FramebufferCreateInfo{
			sType           = .FRAMEBUFFER_CREATE_INFO,
			renderPass      = sc.render_pass,
			attachmentCount = attachment_count,
			pAttachments    = &attachment_buf[0],
			width           = sc.extent.width,
			height          = sc.extent.height,
			layers          = 1,
		}

		result := vk.CreateFramebuffer(dev.device, &fb_info, nil, &sc.framebuffers[i])
		if result != .SUCCESS {
			log.errorf("gpu/core: failed to create framebuffer %d: %v", i, result)
			return false
		}
	}

	return true
}

@(private)
choose_surface_format :: proc(device: vk.PhysicalDevice, surface: vk.SurfaceKHR) -> vk.SurfaceFormatKHR {
	count: u32
	vk.GetPhysicalDeviceSurfaceFormatsKHR(device, surface, &count, nil)
	formats := make([]vk.SurfaceFormatKHR, count, context.temp_allocator)
	vk.GetPhysicalDeviceSurfaceFormatsKHR(device, surface, &count, raw_data(formats))

	// Prefer SRGB B8G8R8A8
	for f in formats {
		if f.format == .B8G8R8A8_SRGB && f.colorSpace == .SRGB_NONLINEAR {
			return f
		}
	}

	return formats[0] if len(formats) > 0 else {}
}

@(private)
choose_present_mode :: proc(device: vk.PhysicalDevice, surface: vk.SurfaceKHR) -> vk.PresentModeKHR {
	count: u32
	vk.GetPhysicalDeviceSurfacePresentModesKHR(device, surface, &count, nil)
	modes := make([]vk.PresentModeKHR, count, context.temp_allocator)
	vk.GetPhysicalDeviceSurfacePresentModesKHR(device, surface, &count, raw_data(modes))

	// Prefer mailbox (triple buffering), fall back to FIFO (vsync)
	for m in modes {
		if m == .MAILBOX {
			return m
		}
	}

	return .FIFO
}

@(private)
choose_extent :: proc(caps: vk.SurfaceCapabilitiesKHR, width, height: u32) -> vk.Extent2D {
	if caps.currentExtent.width != max(u32) {
		return caps.currentExtent
	}

	return vk.Extent2D{
		width  = clamp(width, caps.minImageExtent.width, caps.maxImageExtent.width),
		height = clamp(height, caps.minImageExtent.height, caps.maxImageExtent.height),
	}
}