#include <runtime/mem_arena.h>
#include <runtime/config.h>

#include "device.h"

RT_CVAR_I(r_VkPreferMailboxMode, "Prefer the mailbox present mode. [0/1] (Default: 1)", 1);
RT_CVAR_I(r_VkPreferRelaxedMode, "Prefer the relaxed FIFO present mode. [0/1] (Default: 0)", 0);
RT_CVAR_I(r_VkEnableVSync, "Enable VSync. [0/1] (Default: 0)", 0);

rt_result CreateSwapchain(rt_vk_device *dev) {
    rt_temp_arena temp = rtGetTemporaryArena(NULL, 0);
    uint32_t format_count = 0;
    if (vkGetPhysicalDeviceSurfaceFormatsKHR(dev->phys_device, dev->surface, &format_count, NULL) != VK_SUCCESS) {
        rtReturnTemporaryArena(temp);
        return RT_UNKNOWN_ERROR;
    }
    VkSurfaceFormatKHR *formats = RT_ARENA_PUSH_ARRAY(temp.arena, VkSurfaceFormatKHR, format_count);
    if (vkGetPhysicalDeviceSurfaceFormatsKHR(dev->phys_device, dev->surface, &format_count, formats) != VK_SUCCESS) {
        rtReturnTemporaryArena(temp);
        return RT_UNKNOWN_ERROR;
    }

    /* Pick a format */
    VkFormat surface_format = formats[0].format;
    VkColorSpaceKHR color_space = formats[0].colorSpace;
    for (uint32_t i = 0; i < format_count; ++i) {
        if (formats[i].format == VK_FORMAT_B8G8R8A8_SRGB && formats[i].colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR) {
            surface_format = formats[i].format;
            color_space = formats[i].colorSpace;
            break;
        }
    }

    /* Determine the number of images to create */
    uint32_t num_images = dev->max_frames_in_flight + 1;
    VkSurfaceCapabilitiesKHR surface_capabilities;
    vkGetPhysicalDeviceSurfaceCapabilitiesKHR(dev->phys_device, dev->surface, &surface_capabilities);
    if (surface_capabilities.maxImageCount > 0 && surface_capabilities.maxImageCount < dev->max_frames_in_flight) {
        dev->max_frames_in_flight = surface_capabilities.maxImageCount;
        num_images = dev->max_frames_in_flight;
        rtLog("VK", "Limiting number of frames in flight to maximum number of swapchain images %u.", dev->max_frames_in_flight);
    }
    if (surface_capabilities.maxImageCount > 0 && num_images > surface_capabilities.maxImageCount) {
        num_images = surface_capabilities.maxImageCount;
    }


    /* Determine the extent */
    VkExtent2D extent = surface_capabilities.currentExtent;
    if (surface_capabilities.currentExtent.width == UINT32_MAX && surface_capabilities.currentExtent.height == UINT32_MAX) {
        extent = surface_capabilities.maxImageExtent;
    }

    /* Determine the present mode */
    uint32_t present_mode_count = 0;
    if (vkGetPhysicalDeviceSurfacePresentModesKHR(dev->phys_device, dev->surface, &present_mode_count, NULL) != VK_SUCCESS) {
        rtReturnTemporaryArena(temp);
        return RT_UNKNOWN_ERROR;
    }
    VkPresentModeKHR *present_modes = RT_ARENA_PUSH_ARRAY(temp.arena, VkPresentModeKHR, present_mode_count);
    if (vkGetPhysicalDeviceSurfacePresentModesKHR(dev->phys_device, dev->surface, &present_mode_count, present_modes) != VK_SUCCESS) {
        rtReturnTemporaryArena(temp);
        return RT_UNKNOWN_ERROR;
    }

    VkPresentModeKHR present_mode = present_modes[0];
    if (!r_VkEnableVSync.i) {
        if (r_VkPreferMailboxMode.i) {
            rtLog("VK", "r_VkPreferMailboxMode has no effect, if VSync is disabled.");
            r_VkPreferMailboxMode.i = 0;
            rtNotifyCVARChange(&r_VkPreferMailboxMode);
        }
        if (r_VkPreferRelaxedMode.i) {
            rtLog("VK", "r_VkPreferRelaxedMode has no effect, if VSync is disabled.");
            r_VkPreferRelaxedMode.i = 0;
            rtNotifyCVARChange(&r_VkPreferRelaxedMode);
        }

        for (uint32_t i = 0; i < present_mode_count; ++i) {
            if (present_modes[i] == VK_PRESENT_MODE_IMMEDIATE_KHR) {
                present_mode = VK_PRESENT_MODE_IMMEDIATE_KHR;
                break;
            }
        }

        if (present_mode != VK_PRESENT_MODE_IMMEDIATE_KHR) {
            rtLog("VK", "VSync was disabled, but the physical device/surface combination does not support VK_PRESENT_MODE_IMMEDIATE_KHR");
            r_VkEnableVSync.i = 1;
            rtNotifyCVARChange(&r_VkEnableVSync);
        }
    }
    /* NOT an else if, because we fall back to this if disabled vsync is not supported. */
    if (r_VkEnableVSync.i) {
        /* Required to be supported */
        present_mode = VK_PRESENT_MODE_FIFO_KHR;
        if (r_VkPreferMailboxMode.i || r_VkPreferRelaxedMode.i) {
            for (uint32_t i = 0; i < present_mode_count; ++i) {
                if (present_modes[i] == VK_PRESENT_MODE_FIFO_RELAXED_KHR && r_VkPreferRelaxedMode.i) {
                    present_mode = present_modes[i];
                    break;
                }
                if (present_modes[i] == VK_PRESENT_MODE_MAILBOX_KHR && r_VkPreferMailboxMode.i) {
                    present_mode = present_modes[i];
                    break;
                }
            }
        }
    }
    VkSwapchainCreateInfoKHR swapchain_info = {
        .sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
        .pNext = NULL,
        .flags = 0,
        .surface = dev->surface,
        .minImageCount =  num_images,
        .imageFormat = surface_format,
        .imageColorSpace = color_space,
        .imageExtent = extent,
        .imageArrayLayers = 1,
        .imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,
        .imageSharingMode = VK_SHARING_MODE_EXCLUSIVE,
        .pQueueFamilyIndices = NULL,
        .queueFamilyIndexCount = 0,
        .preTransform =  surface_capabilities.currentTransform,
        .compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR,
        .presentMode = present_mode,
        .clipped = VK_TRUE,
        .oldSwapchain = VK_NULL_HANDLE,
    };

    if (dev->graphics_family != dev->present_family) {
        uint32_t queue_families[2];
        queue_families[0] = dev->graphics_family;
        queue_families[1] = dev->present_family;
        swapchain_info.pQueueFamilyIndices = &queue_families[0];
        swapchain_info.queueFamilyIndexCount = 2;
    }

    if (vkCreateSwapchainKHR(dev->device, &swapchain_info, dev->alloc_cb, &dev->swapchain) != VK_SUCCESS) {
        rtReturnTemporaryArena(temp);
        return RT_UNKNOWN_ERROR;
    }

    vkGetSwapchainImagesKHR(dev->device, dev->swapchain, &dev->swapchain_image_count, NULL);
    RT_ASSERT(dev->swapchain_image_count <= RT_ARRAY_COUNT(dev->swapchain_images), "Unexpectedly high number of swapchain images.");
    vkGetSwapchainImagesKHR(dev->device, dev->swapchain, &dev->swapchain_image_count, &dev->swapchain_images[0]);
    for (uint32_t i = 0; i < dev->swapchain_image_count; ++i) {
        VkImageViewCreateInfo view_info = {
            .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
            .image = dev->swapchain_images[i],
            .viewType = VK_IMAGE_VIEW_TYPE_2D,
            .format = surface_format,
            .components = { VK_COMPONENT_SWIZZLE_IDENTITY, VK_COMPONENT_SWIZZLE_IDENTITY,
                VK_COMPONENT_SWIZZLE_IDENTITY, VK_COMPONENT_SWIZZLE_IDENTITY },
            .subresourceRange = { .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
                .baseArrayLayer =  0,.baseMipLevel = 0,
                .layerCount = 1, .levelCount = 1},
        };
        vkCreateImageView(dev->device, &view_info, dev->alloc_cb, &dev->swapchain_image_views[i]);
    }

    return RT_SUCCESS;
}

void DestroySwapchain(rt_vk_device *dev) {
    vkDestroySwapchainKHR(dev->device, dev->swapchain, dev->alloc_cb);
    for (uint32_t i = 0; i < dev->swapchain_image_count; ++i) {
        vkDestroyImageView(dev->device, dev->swapchain_image_views[i], dev->alloc_cb);
    }
}

rt_result RecreateSwapchain(rt_vk_device *dev) {
    /* TODO(Kevin): Do this in a more perfomant way involving oldSwapchain */
    vkDeviceWaitIdle(dev->device);
    DestroySwapchain(dev);
    return CreateSwapchain(dev);
}
