#include "logicaldevice.h"

#include "swapchain.h"
#include <iostream>

constexpr int MAX_FRAMES_IN_FLIGHT = 2;
constexpr float queuePriority = 1.0f;

Engine::LogicalDevice::LogicalDevice(
    PhysicalDevice* physicalDevice,
    Surface* surface,
    const std::vector<const char*>& extensions,
    const std::vector<const char*> layers
):
    phys(physicalDevice),
    vk(),
    graphicsQueue(),
    presentQueue(),
    renderPass(),
    commandPool(),
    commandBuffers(),
    imageAvailableSemaphores(),
    renderFinishedSemaphores(),
    inFlightFences(),
    currentFrame(0),
    framebufferResized(false),
    surface(surface),
    surfaceFormat(surface->chooseSwapSurfaceFormat(phys->swapChainSupport)),
    swapChain(nullptr)
{
    const QueueFamilyIndices& indices = phys->indices;

    createDevice(phys->vk, indices, extensions, layers);
    createQueues(indices.graphicsFamily.value(), indices.presentFamily.value());

    createCommandPool(phys->indices.graphicsFamily.value());
    createCommandBuffers();
    createSyncObjects();
    createRenderPass(surfaceFormat.format);
}

Engine::LogicalDevice::~LogicalDevice() {
    clearSwapChain();

    pipelines.clear();

    vkDestroyRenderPass(vk, renderPass, nullptr);

    for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++) {
        vkDestroySemaphore(vk, renderFinishedSemaphores[i], nullptr);
        vkDestroySemaphore(vk, imageAvailableSemaphores[i], nullptr);
        vkDestroyFence(vk, inFlightFences[i], nullptr);
    }

    vkFreeCommandBuffers(vk, commandPool, MAX_FRAMES_IN_FLIGHT, commandBuffers.data());
    vkDestroyCommandPool(vk, commandPool, nullptr);

    vkDestroyDevice(vk, nullptr);
}

void Engine::LogicalDevice::addProgram(const Program& program) {
    pipelines.emplace_back(program, this);
}

void Engine::LogicalDevice::createSwapChain() {
    clearSwapChain();
    swapChain = new SwapChain(this);
}

void Engine::LogicalDevice::clearSwapChain() {
    if (swapChain != nullptr)
        delete swapChain;

    swapChain = nullptr;
}

VkResult Engine::LogicalDevice::waitIdle() {
    return vkDeviceWaitIdle(vk);
}

VkResult Engine::LogicalDevice::waitForFence(const VkFence& fence) {
    return vkWaitForFences(vk, 1, &fence, VK_TRUE, UINT64_MAX);
}

VkResult Engine::LogicalDevice::resetFence(const VkFence& fence) {
    return vkResetFences(vk, 1, &fence);
}

VkResult Engine::LogicalDevice::queueSubmitGraphics(
    const VkSemaphore& wait,
    const VkCommandBuffer& buffer,
    const VkSemaphore& signal,
    const VkFence& fence
) {
    VkSubmitInfo submitInfo{};
    submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;

    VkPipelineStageFlags waitStages[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT};
    submitInfo.waitSemaphoreCount = 1;
    submitInfo.pWaitSemaphores = &wait;
    submitInfo.pWaitDstStageMask = waitStages;

    submitInfo.commandBufferCount = 1;
    submitInfo.pCommandBuffers = &buffer;

    submitInfo.signalSemaphoreCount = 1;
    submitInfo.pSignalSemaphores = &signal;

    return vkQueueSubmit(graphicsQueue, 1, &submitInfo, fence);
}

VkResult Engine::LogicalDevice::queuePresent(const VkSemaphore& signal, uint32_t index) {
    VkPresentInfoKHR presentInfo{};
    presentInfo.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;

    presentInfo.waitSemaphoreCount = 1;
    presentInfo.pWaitSemaphores = &signal;

    presentInfo.swapchainCount = 1;
    presentInfo.pSwapchains = &swapChain->vk;

    presentInfo.pImageIndices = &index;

    return vkQueuePresentKHR(presentQueue, &presentInfo);
}

void Engine::LogicalDevice::createDevice(
        VkPhysicalDevice physicalDevice,
        const QueueFamilyIndices& indices,
        const std::vector<const char*>& extensions,
        const std::vector<const char*> layers
) {
    std::vector<VkDeviceQueueCreateInfo> queueCreateInfos;
    std::set<uint32_t> uniqueQueueFamilies = {indices.graphicsFamily.value(), indices.presentFamily.value()};

    for (uint32_t queueFamily : uniqueQueueFamilies) {
        VkDeviceQueueCreateInfo queueCreateInfo{};
        queueCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
        queueCreateInfo.queueFamilyIndex = queueFamily;
        queueCreateInfo.queueCount = 1;
        queueCreateInfo.pQueuePriorities = &queuePriority;
        queueCreateInfos.push_back(queueCreateInfo);
    }

    VkPhysicalDeviceFeatures deviceFeatures{};

    VkDeviceCreateInfo createInfo{};
    createInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;

    createInfo.queueCreateInfoCount = static_cast<uint32_t>(queueCreateInfos.size());
    createInfo.pQueueCreateInfos = queueCreateInfos.data();

    createInfo.pEnabledFeatures = &deviceFeatures;

    createInfo.enabledExtensionCount = static_cast<uint32_t>(extensions.size());
    createInfo.ppEnabledExtensionNames = extensions.data();

    createInfo.enabledLayerCount = static_cast<uint32_t>(layers.size());
    createInfo.ppEnabledLayerNames = layers.data();

    if (vkCreateDevice(physicalDevice, &createInfo, nullptr, &vk) != VK_SUCCESS)
        throw std::runtime_error("failed to create logical device!");
}

void Engine::LogicalDevice::createQueues(uint32_t graphicsFamilyIndex, uint32_t presenFamilyIndex) {
    vkGetDeviceQueue(vk, graphicsFamilyIndex, 0, &graphicsQueue);
    vkGetDeviceQueue(vk, presenFamilyIndex, 0, &presentQueue);
}

void Engine::LogicalDevice::createSyncObjects() {
    imageAvailableSemaphores.resize(MAX_FRAMES_IN_FLIGHT);
    renderFinishedSemaphores.resize(MAX_FRAMES_IN_FLIGHT);
    inFlightFences.resize(MAX_FRAMES_IN_FLIGHT);

    VkSemaphoreCreateInfo semaphoreInfo{};
    semaphoreInfo.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;

    VkFenceCreateInfo fenceInfo{};
    fenceInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
    fenceInfo.flags = VK_FENCE_CREATE_SIGNALED_BIT;

    for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++) {
        if (vkCreateSemaphore(vk, &semaphoreInfo, nullptr, &imageAvailableSemaphores[i]) != VK_SUCCESS ||
            vkCreateSemaphore(vk, &semaphoreInfo, nullptr, &renderFinishedSemaphores[i]) != VK_SUCCESS ||
            vkCreateFence(vk, &fenceInfo, nullptr, &inFlightFences[i]) != VK_SUCCESS) {
            throw std::runtime_error("failed to create synchronization objects for a frame!");
        }
    }
}

void Engine::LogicalDevice::createCommandPool(uint32_t queueFamilyIndex) {
    VkCommandPoolCreateInfo poolInfo{};
    poolInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
    poolInfo.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;
    poolInfo.queueFamilyIndex = queueFamilyIndex;

    if (vkCreateCommandPool(vk, &poolInfo, nullptr, &commandPool) != VK_SUCCESS)
        throw std::runtime_error("failed to create command pool!");
}

void Engine::LogicalDevice::createCommandBuffers() {
    commandBuffers.resize(MAX_FRAMES_IN_FLIGHT);

    VkCommandBufferAllocateInfo allocInfo{};
    allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    allocInfo.commandPool = commandPool;
    allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    allocInfo.commandBufferCount = (uint32_t) commandBuffers.size();

    if (vkAllocateCommandBuffers(vk, &allocInfo, commandBuffers.data()) != VK_SUCCESS)
        throw std::runtime_error("failed to allocate command buffers!");
}

void Engine::LogicalDevice::drawFrame() {
    waitForFence(inFlightFences[currentFrame]);

    uint32_t imageIndex;
    VkResult result = vkAcquireNextImageKHR(vk, swapChain->vk, UINT64_MAX, imageAvailableSemaphores[currentFrame], VK_NULL_HANDLE, &imageIndex);

    if (result == VK_ERROR_OUT_OF_DATE_KHR) {
        recreateSwapChain();
        return;
    } else if (result != VK_SUCCESS && result != VK_SUBOPTIMAL_KHR) {
        throw std::runtime_error("failed to acquire swap chain image!");
    }

    resetFence(inFlightFences[currentFrame]);

    if (!pipelines.empty()) {
        vkResetCommandBuffer(commandBuffers[currentFrame], /*VkCommandBufferResetFlagBits*/ 0);
        recordCommandBuffer(commandBuffers[currentFrame], imageIndex);
    } else {
        std::cout << "no pipelines attached, skipping command buffer record" << std::endl;
    }

    if (queueSubmitGraphics(
        imageAvailableSemaphores[currentFrame],
        commandBuffers[currentFrame],
        renderFinishedSemaphores[currentFrame],
        inFlightFences[currentFrame]) != VK_SUCCESS
    )
        throw std::runtime_error("failed to submit draw command buffer!");

    result = queuePresent(renderFinishedSemaphores[currentFrame], imageIndex);
    if (result == VK_ERROR_OUT_OF_DATE_KHR || result == VK_SUBOPTIMAL_KHR || framebufferResized) {
        framebufferResized = false;
        recreateSwapChain();
    } else if (result != VK_SUCCESS) {
        throw std::runtime_error("failed to present swap chain image!");
    }

    currentFrame = (currentFrame + 1) % MAX_FRAMES_IN_FLIGHT;
}

void Engine::LogicalDevice::setResized() {
    framebufferResized = true;
}

VkExtent2D Engine::LogicalDevice::chooseSwapExtent() const {
    return surface->chooseSwapExtent(phys->swapChainSupport.capabilities);
}

VkResult Engine::LogicalDevice::createVkSwapChain(const VkExtent2D& extent, VkSwapchainKHR& out) {
    const SwapChainSupportDetails& swapChainSupport = phys->swapChainSupport;

    uint32_t imageCount = swapChainSupport.capabilities.minImageCount + 1;
    if (swapChainSupport.capabilities.maxImageCount > 0 && imageCount > swapChainSupport.capabilities.maxImageCount)
        imageCount = swapChainSupport.capabilities.maxImageCount;

    VkSwapchainCreateInfoKHR createInfo{};
    createInfo.sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR;
    createInfo.surface = surface->vk;

    createInfo.imageFormat = surfaceFormat.format;
    createInfo.imageColorSpace = surfaceFormat.colorSpace;
    createInfo.imageArrayLayers = 1;
    createInfo.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;

    createInfo.imageExtent = extent;
    createInfo.minImageCount = imageCount;

    const QueueFamilyIndices& indices = phys->indices;
    uint32_t queueFamilyIndices[] = {indices.graphicsFamily.value(), indices.presentFamily.value()};

    if (indices.graphicsFamily != indices.presentFamily) {
        createInfo.imageSharingMode = VK_SHARING_MODE_CONCURRENT;
        createInfo.queueFamilyIndexCount = 2;
        createInfo.pQueueFamilyIndices = queueFamilyIndices;
    } else {
        createInfo.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
    }

    createInfo.preTransform = swapChainSupport.capabilities.currentTransform;
    createInfo.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;
    createInfo.presentMode = chooseSwapPresentMode(swapChainSupport.presentModes);
    createInfo.clipped = VK_TRUE;

    return vkCreateSwapchainKHR(vk, &createInfo, nullptr, &out);
}

void Engine::LogicalDevice::destroyVkSwapChain(VkSwapchainKHR& swapChain) {
    vkDestroySwapchainKHR(vk, swapChain, nullptr);
}

VkResult Engine::LogicalDevice::getVkSwapChainImages(VkSwapchainKHR& swapChain, std::vector<VkImage>& out) {
    uint32_t imageCount;
    VkResult result =  vkGetSwapchainImagesKHR(vk, swapChain, &imageCount, nullptr);
    if (result != VK_SUCCESS)
        return result;

    out.resize(imageCount);
    return vkGetSwapchainImagesKHR(vk, swapChain, &imageCount, out.data());
}

VkResult Engine::LogicalDevice::createVkImageView(const VkImage& image, VkImageView& out) {
    VkImageViewCreateInfo createInfo{};
    createInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
    createInfo.image = image;
    createInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
    createInfo.format = surfaceFormat.format;
    createInfo.components.r = VK_COMPONENT_SWIZZLE_IDENTITY;
    createInfo.components.g = VK_COMPONENT_SWIZZLE_IDENTITY;
    createInfo.components.b = VK_COMPONENT_SWIZZLE_IDENTITY;
    createInfo.components.a = VK_COMPONENT_SWIZZLE_IDENTITY;
    createInfo.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    createInfo.subresourceRange.baseMipLevel = 0;
    createInfo.subresourceRange.levelCount = 1;
    createInfo.subresourceRange.baseArrayLayer = 0;
    createInfo.subresourceRange.layerCount = 1;

    return vkCreateImageView(vk, &createInfo, nullptr, &out);
}

void Engine::LogicalDevice::destroyVkImageView(VkImageView& view) {
    vkDestroyImageView(vk, view, nullptr);
}

VkResult Engine::LogicalDevice::createVkFrameBuffer(const VkImageView& imageView, const VkExtent2D& extent, VkFramebuffer& out) {
    VkImageView attachments[] = {imageView};

    VkFramebufferCreateInfo framebufferInfo{};
    framebufferInfo.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;
    framebufferInfo.renderPass = renderPass;
    framebufferInfo.attachmentCount = 1;
    framebufferInfo.pAttachments = attachments;
    framebufferInfo.width = extent.width;
    framebufferInfo.height = extent.height;
    framebufferInfo.layers = 1;

    return vkCreateFramebuffer(vk, &framebufferInfo, nullptr, &out);
}

void Engine::LogicalDevice::destroyVkFrameBuffer(VkFramebuffer& buffer) {
    vkDestroyFramebuffer(vk, buffer, nullptr);
}

void Engine::LogicalDevice::createRenderPass(VkFormat format) {
    VkAttachmentDescription colorAttachment{};
    colorAttachment.format = format;
    colorAttachment.samples = VK_SAMPLE_COUNT_1_BIT;
    colorAttachment.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
    colorAttachment.storeOp = VK_ATTACHMENT_STORE_OP_STORE;
    colorAttachment.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
    colorAttachment.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
    colorAttachment.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
    colorAttachment.finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;

    VkAttachmentReference colorAttachmentRef{};
    colorAttachmentRef.attachment = 0;
    colorAttachmentRef.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;

    VkSubpassDescription subpass{};
    subpass.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
    subpass.colorAttachmentCount = 1;
    subpass.pColorAttachments = &colorAttachmentRef;

    VkSubpassDependency dependency{};
    dependency.srcSubpass = VK_SUBPASS_EXTERNAL;
    dependency.dstSubpass = 0;
    dependency.srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
    dependency.srcAccessMask = 0;
    dependency.dstStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
    dependency.dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;

    VkRenderPassCreateInfo renderPassInfo{};
    renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO;
    renderPassInfo.attachmentCount = 1;
    renderPassInfo.pAttachments = &colorAttachment;
    renderPassInfo.subpassCount = 1;
    renderPassInfo.pSubpasses = &subpass;
    renderPassInfo.dependencyCount = 1;
    renderPassInfo.pDependencies = &dependency;

    if (vkCreateRenderPass(vk, &renderPassInfo, nullptr, &renderPass) != VK_SUCCESS)
        throw std::runtime_error("failed to create render pass!");
}

void Engine::LogicalDevice::recordCommandBuffer(VkCommandBuffer commandBuffer, uint32_t imageIndex) {
    VkExtent2D extent = swapChain->getExtent();
    VkCommandBufferBeginInfo beginInfo{};
    beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;

    if (vkBeginCommandBuffer(commandBuffer, &beginInfo) != VK_SUCCESS)
        throw std::runtime_error("failed to begin recording command buffer!");

    VkRenderPassBeginInfo renderPassInfo{};
    renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
    renderPassInfo.renderPass = renderPass;
    renderPassInfo.framebuffer = swapChain->getFrameBuffer(imageIndex);
    renderPassInfo.renderArea.offset = {0, 0};
    renderPassInfo.renderArea.extent = extent;

    VkClearValue clearColor = {{{0.0f, 0.0f, 0.0f, 1.0f}}};
    renderPassInfo.clearValueCount = 1;
    renderPassInfo.pClearValues = &clearColor;

    vkCmdBeginRenderPass(commandBuffer, &renderPassInfo, VK_SUBPASS_CONTENTS_INLINE);

    for (const Pipeline& pipeline: pipelines) {
        pipeline.bind(commandBuffer);

        VkViewport viewport{};
        viewport.x = 0.0f;
        viewport.y = 0.0f;
        viewport.width = (float) extent.width;
        viewport.height = (float) extent.height;
        viewport.minDepth = 0.0f;
        viewport.maxDepth = 1.0f;
        vkCmdSetViewport(commandBuffer, 0, 1, &viewport);

        VkRect2D scissor{};
        scissor.offset = {0, 0};
        scissor.extent = extent;
        vkCmdSetScissor(commandBuffer, 0, 1, &scissor);

        vkCmdDraw(commandBuffer, 3, 1, 0, 0);
    }

    vkCmdEndRenderPass(commandBuffer);

    if (vkEndCommandBuffer(commandBuffer) != VK_SUCCESS)
        throw std::runtime_error("failed to record command buffer!");
}

void Engine::LogicalDevice::recreateSwapChain() {
    surface->waitForResize();
    waitIdle();
    phys->recreateSwapChainSupportDetails(surface);
    createSwapChain();
}

VkResult Engine::LogicalDevice::createPipelineLayout(const VkPipelineLayoutCreateInfo& info, VkPipelineLayout& out) {
    return vkCreatePipelineLayout(vk, &info, nullptr, &out);
}

VkResult Engine::LogicalDevice::createPipeline(VkGraphicsPipelineCreateInfo& info, VkPipeline& out) {
    info.renderPass = renderPass;
    return vkCreateGraphicsPipelines(vk, VK_NULL_HANDLE, 1, &info, nullptr, &out);
}

VkResult Engine::LogicalDevice::createShaderModule(const VkShaderModuleCreateInfo& info, VkShaderModule& out) {
    return vkCreateShaderModule(vk, &info, nullptr, &out);
}

void Engine::LogicalDevice::destroyPipeline(VkPipeline& pipline) {
    vkDestroyPipeline(vk, pipline, nullptr);
}

void Engine::LogicalDevice::destroyPipelineLayout(VkPipelineLayout& layout) {
    vkDestroyPipelineLayout(vk, layout, nullptr);
}

void Engine::LogicalDevice::destroyShaderModule(VkShaderModule& module) {
    vkDestroyShaderModule(vk, module, nullptr);
}