/* SPDX-License-Identifier:BSD-3-Clause */
#include "swapchain.h"
#include "../../core/log.h"
#include <stdlib.h>
#include <string.h>

#define RGFW_IMPORT
#include "../../rgfw.h"

static VkSurfaceFormatKHR select_surface_format(VkPhysicalDevice device, VkSurfaceKHR surface)
{
	u32 format_count = 0;
	vkGetPhysicalDeviceSurfaceFormatsKHR(device, surface, &format_count, NULL);

	VkSurfaceFormatKHR *formats = malloc(sizeof(VkSurfaceFormatKHR) * format_count);
	vkGetPhysicalDeviceSurfaceFormatsKHR(device, surface, &format_count, formats);

	VkSurfaceFormatKHR selected = formats[0];
	for (u32 i = 0; i < format_count; i++) {
		if (formats[i].format == VK_FORMAT_B8G8R8A8_SRGB &&
		    formats[i].colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR) {
			selected = formats[i];
			break;
		}
	}

	free(formats);
	return selected;
}

static VkPresentModeKHR select_present_mode(VkPhysicalDevice device, VkSurfaceKHR surface)
{
	u32 mode_count = 0;
	vkGetPhysicalDeviceSurfacePresentModesKHR(device, surface, &mode_count, NULL);

	VkPresentModeKHR *modes = malloc(sizeof(VkPresentModeKHR) * mode_count);
	vkGetPhysicalDeviceSurfacePresentModesKHR(device, surface, &mode_count, modes);

	VkPresentModeKHR selected = VK_PRESENT_MODE_FIFO_KHR;
	for (u32 i = 0; i < mode_count; i++) {
		if (modes[i] == VK_PRESENT_MODE_MAILBOX_KHR) {
			selected = VK_PRESENT_MODE_MAILBOX_KHR;
			break;
		}
	}

	free(modes);
	return selected;
}

static VkExtent2D select_extent(VkSurfaceCapabilitiesKHR *caps, RGFW_window *window)
{
	if (caps->currentExtent.width != 0xFFFFFFFF) {
		return caps->currentExtent;
	}

	i32 w, h;
	RGFW_window_getSize(window, &w, &h);

	VkExtent2D extent;
	extent.width = (u32)w;
	extent.height = (u32)h;

	if (extent.width < caps->minImageExtent.width)
		extent.width = caps->minImageExtent.width;
	if (extent.width > caps->maxImageExtent.width)
		extent.width = caps->maxImageExtent.width;
	if (extent.height < caps->minImageExtent.height)
		extent.height = caps->minImageExtent.height;
	if (extent.height > caps->maxImageExtent.height)
		extent.height = caps->maxImageExtent.height;

	return extent;
}

void vk_swapchain_init(struct renderer_context *context)
{
	VkSurfaceCapabilitiesKHR caps;
	vkGetPhysicalDeviceSurfaceCapabilitiesKHR(context->physical_device, context->surface, &caps);

	VkSurfaceFormatKHR format = select_surface_format(context->physical_device, context->surface);
	VkPresentModeKHR present_mode = select_present_mode(context->physical_device, context->surface);
	VkExtent2D extent = select_extent(&caps, context->window);

	u32 image_count = caps.minImageCount + 1;
	if (caps.maxImageCount > 0 && image_count > caps.maxImageCount) {
		image_count = caps.maxImageCount;
	}

	VkSwapchainCreateInfoKHR create_info = {
		.sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
		.surface = context->surface,
		.minImageCount = image_count,
		.imageFormat = format.format,
		.imageColorSpace = format.colorSpace,
		.imageExtent = extent,
		.imageArrayLayers = 1,
		.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,
		.preTransform = caps.currentTransform,
		.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR,
		.presentMode = present_mode,
		.clipped = VK_TRUE,
		.oldSwapchain = VK_NULL_HANDLE
	};

	u32 queue_indices[] = { context->queue_indices.graphics, context->queue_indices.present };
	if (context->queue_indices.graphics != context->queue_indices.present) {
		create_info.imageSharingMode = VK_SHARING_MODE_CONCURRENT;
		create_info.queueFamilyIndexCount = 2;
		create_info.pQueueFamilyIndices = queue_indices;
	} else {
		create_info.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
	}

	if (vkCreateSwapchainKHR(context->device, &create_info, NULL, &context->swapchain.handle) != VK_SUCCESS) {
		fatal("Can't create swapchain.\n");
	}

	context->swapchain.format = format.format;
	context->swapchain.extent = extent;

	vkGetSwapchainImagesKHR(context->device, context->swapchain.handle, &context->swapchain.image_count, NULL);
	context->swapchain.images = malloc(sizeof(VkImage) * context->swapchain.image_count);
	vkGetSwapchainImagesKHR(context->device, context->swapchain.handle, &context->swapchain.image_count, context->swapchain.images);

	context->swapchain.image_views = malloc(sizeof(VkImageView) * context->swapchain.image_count);
	for (u32 i = 0; i < context->swapchain.image_count; i++) {
		VkImageViewCreateInfo view_info = {
			.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
			.image = context->swapchain.images[i],
			.viewType = VK_IMAGE_VIEW_TYPE_2D,
			.format = context->swapchain.format,
			.components.r = VK_COMPONENT_SWIZZLE_IDENTITY,
			.components.g = VK_COMPONENT_SWIZZLE_IDENTITY,
			.components.b = VK_COMPONENT_SWIZZLE_IDENTITY,
			.components.a = VK_COMPONENT_SWIZZLE_IDENTITY,
			.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
			.subresourceRange.baseMipLevel = 0,
			.subresourceRange.levelCount = 1,
			.subresourceRange.baseArrayLayer = 0,
			.subresourceRange.layerCount = 1
		};

		if (vkCreateImageView(context->device, &view_info, NULL, &context->swapchain.image_views[i]) != VK_SUCCESS) {
			fatal("Can't create image view.\n");
		}
	}

	log_info("Swapchain created.\n");
}

void vk_swapchain_deinit(struct renderer_context *context)
{
	for (u32 i = 0; i < context->swapchain.image_count; i++) {
		vkDestroyImageView(context->device, context->swapchain.image_views[i], NULL);
	}
	free(context->swapchain.image_views);
	free(context->swapchain.images);
	vkDestroySwapchainKHR(context->device, context->swapchain.handle, NULL);
}