#include "vulkan_swapchain.h"

#include "GLFW/glfw3.h"
#include "imgui.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);
    }

    VkSurfaceCapabilitiesKHR surface_capabilities;
    CHECK_VULKAN(vkGetPhysicalDeviceSurfaceCapabilitiesKHR(
        device.physical_device,
        surface,
        &surface_capabilities));
    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 = {new_width, new_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,
        VmaAllocationCreateInfo {
            .flags = VMA_ALLOCATION_CREATE_DEDICATED_MEMORY_BIT,
            .usage = VMA_MEMORY_USAGE_AUTO,
        });
}

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

    std::cerr << "Creating surface" << std::endl;
    // 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 semaphores
    for (uint32_t i = 0; i < SWAPCHAIN_IMAGE_COUNT; i++) {
        VkSemaphoreCreateInfo semaphore_info = {
            .sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO,
        };
        CHECK_VULKAN(vkCreateSemaphore(
            device.device,
            &semaphore_info,
            nullptr,
            &image_available_semaphores[i]));
        CHECK_VULKAN(vkCreateSemaphore(
            device.device,
            &semaphore_info,
            nullptr,
            &render_finished_semaphores[i]));
    }

    // 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_context = ImGui::CreateContext();
    ImGui::SetCurrentContext(imgui_context);

    ImGuiIO &io = ImGui::GetIO();
    io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard;
    io.FontDefault = io.Fonts->AddFontFromFileTTF(
        "assets/Roboto-Regular.ttf",
        16.0f);

    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,
        .RenderPass = render_pass,
        .MinImageCount = SWAPCHAIN_IMAGE_COUNT,
        .ImageCount = SWAPCHAIN_IMAGE_COUNT,
        .MSAASamples = VK_SAMPLE_COUNT_1_BIT,
    };
    ImGui_ImplVulkan_Init(&init_info);
}

void Swapchain::start_frame() {
    VkCommandBufferBeginInfo begin_info = {
        .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
        .flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT,
    };
    CHECK_VULKAN(vkBeginCommandBuffer(cmd_buf.buffer, &begin_info));
}

void Swapchain::display(Texture2D& texture) {
    std::cerr << "Displaying" << std::endl;
    VkResult present_result = vkAcquireNextImageKHR(
        device.device,
        swapchain,
        UINT64_MAX,
        image_available_semaphores[semaphore_index],
        VK_NULL_HANDLE,
        &frame_index);
    switch (present_result) {
        case VK_ERROR_OUT_OF_DATE_KHR:
        case VK_SUBOPTIMAL_KHR:
            needs_recreate = true;
            return;
        case VK_SUCCESS:
            break;
        default:
            CHECK_VULKAN(present_result);
    }

    std::cerr << "starting image layout transition" << std::endl;
    // Command buffer begins at the very start of the frame
    // Blit the texture to the swapchain image
    {
        // swapchain image layout, undefined -> transfer dst
        VkImageMemoryBarrier swapchain_image_barrier = {
            .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
            .srcAccessMask = 0,
            .dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT,
            .oldLayout = VK_IMAGE_LAYOUT_UNDEFINED,
            .newLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
            .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .image = swapchain_images[frame_index],
            .subresourceRange = {
                .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
                .levelCount = 1,
                .layerCount = 1,
            },
        };
        // src image layout, whatever -> transfer src
        VkImageMemoryBarrier src_image_barrier = {
            .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
            .srcAccessMask = 0,
            .dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT,
            .oldLayout = texture->layout,
            .newLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
            .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
            .image = texture->image,
            .subresourceRange = {
                .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
                .levelCount = 1,
                .layerCount = 1,
            },
        };
        VkImageMemoryBarrier barriers[2] = {swapchain_image_barrier, src_image_barrier};
        std::cerr << "starting image layout transition of 2" << std::endl;
        vkCmdPipelineBarrier(
            cmd_buf.buffer,
            VK_PIPELINE_STAGE_ALL_COMMANDS_BIT,
            VK_PIPELINE_STAGE_TRANSFER_BIT,
            0,
            0, nullptr,
            0, nullptr,
            2, barriers);

        // Blit the texture to the swapchain image
        VkImageBlit blit = {
            .srcSubresource = {
                .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
                .mipLevel = 0,
                .baseArrayLayer = 0,
                .layerCount = 1,
            },
            .srcOffsets = {
                {0, 0, 0},
                {static_cast<int32_t>(texture->extent.width), static_cast<int32_t>(texture->extent.height), 1},
            },
            .dstSubresource = {
                .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
                .mipLevel = 0,
                .baseArrayLayer = 0,
                .layerCount = 1,
            },
            .dstOffsets = {
                {0, 0, 0},
                {static_cast<int32_t>(width), static_cast<int32_t>(height), 1},
            },
        };
        std::cerr << "Starting blit" << std::endl;
        vkCmdBlitImage(
            cmd_buf.buffer,
            texture->image,
            VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
            swapchain_images[frame_index],
            VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
            1,
            &blit,
            VK_FILTER_NEAREST);

        // swapchain image layout, transfer dst -> color attachment
        // re-use the swapchain_image_barrier
        swapchain_image_barrier.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
        swapchain_image_barrier.newLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
        swapchain_image_barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
        swapchain_image_barrier.dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_READ_BIT | 
                                                VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
        vkCmdPipelineBarrier(
            cmd_buf.buffer,
            VK_PIPELINE_STAGE_ALL_COMMANDS_BIT,
            VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
            0,
            0, nullptr,
            0, nullptr,
            1, &swapchain_image_barrier);
    }

    std::cerr << "starting render pass" << std::endl;
    // Render
    VkRenderPassBeginInfo render_pass_begin_info = {
        .sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO,
        .renderPass = render_pass,
        .framebuffer = framebuffers[frame_index],
        .renderArea = {
            .offset = {0, 0},
            .extent = {width, height},
        },
        .clearValueCount = 0,
    };
    vkCmdBeginRenderPass(
        cmd_buf.buffer,
        &render_pass_begin_info,
        VK_SUBPASS_CONTENTS_INLINE);
    std::cerr << "rendering imgui" << std::endl;
    ImGui_ImplVulkan_RenderDrawData(ImGui::GetDrawData(), cmd_buf.buffer);
    vkCmdEndRenderPass(cmd_buf.buffer);

    // End command buffer and submit
    CHECK_VULKAN(vkEndCommandBuffer(cmd_buf.buffer));
    {
        VkPipelineStageFlags wait_stage = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
        VkSubmitInfo submit_info = {
            .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO,
            .waitSemaphoreCount = 1,
            .pWaitSemaphores = &image_available_semaphores[semaphore_index],
            .pWaitDstStageMask = &wait_stage,
            .commandBufferCount = 1,
            .pCommandBuffers = &cmd_buf.buffer,
            .signalSemaphoreCount = 1,
            .pSignalSemaphores = &render_finished_semaphores[semaphore_index],
        };
        CHECK_VULKAN(vkQueueSubmit(
            cmd_buf.queue,
            1,
            &submit_info,
            VK_NULL_HANDLE));
    }

    // Present the image to the swapchain
    VkPresentInfoKHR present_info = {
        .sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR,
        .waitSemaphoreCount = 1,
        .pWaitSemaphores = &render_finished_semaphores[semaphore_index],
        .swapchainCount = 1,
        .pSwapchains = &swapchain,
        .pImageIndices = &frame_index,
    };
    present_result = vkQueuePresentKHR(device.graphics_queue, &present_info);

    switch (present_result) {
        case VK_ERROR_OUT_OF_DATE_KHR:
        case VK_SUBOPTIMAL_KHR:
            needs_recreate = true;
        case VK_SUCCESS:
            break;
        default:
            CHECK_VULKAN(present_result);
    }

    // Rotate the semaphore index
    semaphore_index = (semaphore_index + 1) % SWAPCHAIN_IMAGE_COUNT;
}

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);
        vkDestroySemaphore(device.device, image_available_semaphores[i], nullptr);
        vkDestroySemaphore(device.device, render_finished_semaphores[i], nullptr);
    }

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

    upload_texture.reset();
}

} // namespace iris