#include "vulkan_swapchain.h"

#include "GLFW/glfw3.h"
#include "imgui_impl_glfw.h"
#include "imgui_impl_vulkan.h"
#include "vulkan/vulkan_core.h"

#include <cstdint>
#include <cstdlib>
#include <vector>
#include <iostream>
#include <string>
#include <vk_mem_alloc.h>

std::vector<std::string> get_glfw_instance_extensions() {
    uint32_t extension_count = 0;
    const char **extensions = glfwGetRequiredInstanceExtensions(&extension_count);

    std::vector<std::string> result;
    result.reserve(extension_count);
    for (uint32_t i = 0; i < extension_count; i++) {
        result.push_back(extensions[i]);
    }
    return result;
}

namespace iris {

void Swapchain::resize(uint32_t new_width, uint32_t new_height) {
    width = new_width;
    height = new_height;

    // wait for the device to finish
    CHECK_VULKAN(vkDeviceWaitIdle(device.device));

    // destroy old resources
    if (swapchain != VK_NULL_HANDLE) {
        for (uint32_t i = 0; i < SWAPCHAIN_IMAGE_COUNT; i++) {
            vkDestroyFramebuffer(device.device, framebuffers[i], nullptr);
            vkDestroyImageView(device.device, swapchain_image_views[i], nullptr);
        }
        vkDestroySwapchainKHR(device.device, swapchain, nullptr);
    }

    VkSwapchainCreateInfoKHR swapchain_create_info = {
        .sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
        .surface = surface,
        .minImageCount = SWAPCHAIN_IMAGE_COUNT,
        .imageFormat = VK_FORMAT_B8G8R8A8_UNORM,
        .imageColorSpace = VK_COLOR_SPACE_SRGB_NONLINEAR_KHR,
        .imageExtent = {width, height},
        .imageArrayLayers = 1,
        .imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT,
        .imageSharingMode = VK_SHARING_MODE_EXCLUSIVE,
        .preTransform = VK_SURFACE_TRANSFORM_IDENTITY_BIT_KHR,
        .compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR,
        .presentMode = VK_PRESENT_MODE_IMMEDIATE_KHR,
        .clipped = VK_TRUE,
        .oldSwapchain = VK_NULL_HANDLE,
    };
    CHECK_VULKAN(vkCreateSwapchainKHR(
        device.device,
        &swapchain_create_info,
        nullptr,
        &swapchain));

    // images
    uint32_t image_count = 0;
    CHECK_VULKAN(vkGetSwapchainImagesKHR(
        device.device,
        swapchain,
        &image_count,
        nullptr));
    if (image_count > SWAPCHAIN_IMAGE_COUNT) {
        // TODO throw an exception
        std::cerr << "Swapchain image count is greater than expected" << std::endl;
        abort();
    }
    CHECK_VULKAN(vkGetSwapchainImagesKHR(
        device.device,
        swapchain,
        &image_count,
        swapchain_images));

    // image views
    for (uint32_t i = 0; i < image_count; i++) {
        VkImageViewCreateInfo image_view_create_info = {
            .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
            .image = swapchain_images[i],
            .viewType = VK_IMAGE_VIEW_TYPE_2D,
            .format = VK_FORMAT_B8G8R8A8_UNORM,
            .components = {
                .r = VK_COMPONENT_SWIZZLE_IDENTITY,
                .g = VK_COMPONENT_SWIZZLE_IDENTITY,
                .b = VK_COMPONENT_SWIZZLE_IDENTITY,
                .a = VK_COMPONENT_SWIZZLE_IDENTITY,
            },
            .subresourceRange = {
                .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
                .baseMipLevel = 0,
                .levelCount = 1,
                .baseArrayLayer = 0,
                .layerCount = 1,
            },
        };

        CHECK_VULKAN(vkCreateImageView(
            device.device,
            &image_view_create_info,
            nullptr,
            &swapchain_image_views[i]));
    }

    // framebuffers
    for (uint32_t i = 0; i < image_count; i++) {
        VkFramebufferCreateInfo frame_buffer_create_info = {
            .sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO,
            .renderPass = render_pass,
            .attachmentCount = 1,
            .pAttachments = &swapchain_image_views[i],
            .width = width,
            .height = height,
            .layers = 1,
        };
        CHECK_VULKAN(vkCreateFramebuffer(
            device.device,
            &frame_buffer_create_info,
            nullptr,
            &framebuffers[i]));
    }

    upload_texture = device.create_texture(
        {width, height},
        VK_FORMAT_B8G8R8A8_UNORM,
        VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT,
        VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE);
}

Swapchain::Swapchain(
    GLFWwindow *window,
    iris::Device device) : device(device), window(window)
{
    if (!glfwVulkanSupported()) {
        std::cerr << "GLFW failed to find Vulkan support" << std::endl;
        // TODO throw an exception
        abort();
    }

    // Create the surface
    CHECK_VULKAN(glfwCreateWindowSurface(
        device.instance,
        window,
        VK_NULL_HANDLE,
        &surface));

    // Optionally check surface capabilities
    {
        VkBool32 supported = VK_FALSE;
        CHECK_VULKAN(vkGetPhysicalDeviceSurfaceSupportKHR(
            device.physical_device,
            device.main_queue_family_index,
            surface,
            &supported));
        if (supported != VK_TRUE) {
            // TODO throw an exception
            std::cerr << "Surface does not support presentation" << std::endl;
            abort();
        }  

        VkSurfaceCapabilitiesKHR surface_capabilities;
        CHECK_VULKAN(vkGetPhysicalDeviceSurfaceCapabilitiesKHR(
            device.physical_device,
            surface,
            &surface_capabilities));

        if (surface_capabilities.maxImageCount < SWAPCHAIN_IMAGE_COUNT) {
            // TODO throw an exception
            std::cerr << "Surface does not support enough images" << std::endl;
            abort();
        }
    }
    
    // Create render pass
    VkAttachmentDescription2 attachment = {
        .sType = VK_STRUCTURE_TYPE_ATTACHMENT_DESCRIPTION_2,
        .format = VK_FORMAT_B8G8R8A8_UNORM,
        .samples = VK_SAMPLE_COUNT_1_BIT,
        .loadOp = VK_ATTACHMENT_LOAD_OP_LOAD,
        .storeOp = VK_ATTACHMENT_STORE_OP_STORE,
        .stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE,
        .stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE,
        .initialLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
        .finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
    };
    VkAttachmentReference2 color_attachment = {
        .sType = VK_STRUCTURE_TYPE_ATTACHMENT_REFERENCE_2,
        .attachment = 0,
        .layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
    };
    VkSubpassDescription2 subpass = {
        .sType = VK_STRUCTURE_TYPE_SUBPASS_DESCRIPTION_2,
        .pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS,
        .colorAttachmentCount = 1,
        .pColorAttachments = &color_attachment,
    };
    VkSubpassDependency2 dependencies[2] = {
        {
            .sType = VK_STRUCTURE_TYPE_SUBPASS_DEPENDENCY_2,
            .srcSubpass = VK_SUBPASS_EXTERNAL,
            .dstSubpass = 0,
            .srcStageMask = VK_PIPELINE_STAGE_TRANSFER_BIT | VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
            .dstStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
            .srcAccessMask = VK_ACCESS_MEMORY_WRITE_BIT,
            .dstAccessMask = VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT,
        },
        {
            .sType = VK_STRUCTURE_TYPE_SUBPASS_DEPENDENCY_2,
            .srcSubpass = 0,
            .dstSubpass = VK_SUBPASS_EXTERNAL,
            .srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
            .dstStageMask = VK_PIPELINE_STAGE_TRANSFER_BIT | VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
            .srcAccessMask = VK_ACCESS_MEMORY_WRITE_BIT,
            .dstAccessMask = VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT,
        },
    };

    VkRenderPassCreateInfo2 render_pass_create_info = {
        .sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO_2,
        .attachmentCount = 1,
        .pAttachments = &attachment,
        .subpassCount = 1,
        .pSubpasses = &subpass,
        .dependencyCount = 2,
        .pDependencies = dependencies,
    };
    CHECK_VULKAN(vkCreateRenderPass2(
        device.device,
        &render_pass_create_info,
        nullptr,
        &render_pass));

    // Create swapchain and image/imageview/framebuffers
    this->resize(width, height);

    // Create descriptor pool
    VkDescriptorPoolSize pool_size = {
        .type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
        .descriptorCount = 1,
    };
    VkDescriptorPoolCreateInfo pool_info = {
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
        .maxSets = 1,
        .poolSizeCount = 1,
        .pPoolSizes = &pool_size,
    };
    CHECK_VULKAN(vkCreateDescriptorPool(
        device.device,
        &pool_info,
        nullptr,
        &descriptor_pool));

    // initialize ImGui
    ImGui_ImplGlfw_InitForVulkan(window, true);
    ImGui_ImplVulkan_InitInfo init_info = {
        .Instance = device.instance,
        .PhysicalDevice = device.physical_device,
        .Device = device.device,
        .QueueFamily = device.main_queue_family_index,
        .Queue = device.graphics_queue,
        .DescriptorPool = descriptor_pool,
        .MinImageCount = SWAPCHAIN_IMAGE_COUNT,
        .ImageCount = SWAPCHAIN_IMAGE_COUNT,
        .MSAASamples = VK_SAMPLE_COUNT_1_BIT,

        // optional 
        .CheckVkResultFn = [](const VkResult err) { CHECK_VULKAN(err); },
    };
    ImGui_ImplVulkan_Init(&init_info);
}

Swapchain::~Swapchain() {
    ImGui_ImplVulkan_Shutdown();

    for (uint32_t i = 0; i < SWAPCHAIN_IMAGE_COUNT; i++) {
        vkDestroyFramebuffer(device.device, framebuffers[i], nullptr);
        vkDestroyImageView(device.device, swapchain_image_views[i], nullptr);
    }

    vkDestroyDescriptorPool(device.device, descriptor_pool, nullptr);
    vkDestroySwapchainKHR(device.device, swapchain, nullptr);
    vkDestroyRenderPass(device.device, render_pass, nullptr);
    vkDestroySurfaceKHR(device.instance, surface, nullptr);
}

} // namespace iris