#include "RenderPipeline.hpp"
#include "TileLayout.hpp"

#include <array>
#include <cstring>
#include <fstream>
#include <iostream>
#include <stdexcept>

std::vector<char> RenderPipeline::readFile(
    const std::string& filename)
{
    std::ifstream file(
        filename,
        std::ios::ate | std::ios::binary
    );

    if (!file.is_open())
    {
        throw std::runtime_error(
            "Failed to open file: " + filename
        );
    }

    size_t fileSize =
        static_cast<size_t>(file.tellg());

    std::vector<char> buffer(fileSize);
    file.seekg(0);
    file.read(buffer.data(), fileSize);
    file.close();

    return buffer;
}

VkShaderModule RenderPipeline::createShaderModule(
    const std::vector<char>& code)
{
    VkShaderModuleCreateInfo createInfo{};
    createInfo.sType =
        VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
    createInfo.codeSize = code.size();
    createInfo.pCode =
        reinterpret_cast<const uint32_t*>(
            code.data()
        );

    VkShaderModule shaderModule;

    if (vkCreateShaderModule(
            m_device,
            &createInfo,
            nullptr,
            &shaderModule) != VK_SUCCESS)
    {
        throw std::runtime_error(
            "Failed to create shader module"
        );
    }

    return shaderModule;
}

RenderPipeline::RenderPipeline(
    VkDevice device,
    VkFormat swapchainImageFormat)
    : m_device(device)
{
    // Render pass
    VkAttachmentDescription colorAttachment{};
    colorAttachment.format = swapchainImageFormat;
    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;

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

    if (vkCreateRenderPass(
            m_device,
            &renderPassInfo,
            nullptr,
            &m_renderPass) != VK_SUCCESS)
    {
        throw std::runtime_error(
            "Failed to create render pass"
        );
    }

    std::cout << "Render pass created" << std::endl;

    // Shaders
    auto vertShaderCode =
        readFile(
            "../shaders/bin/textured.vert.spv"
        );
    auto fragShaderCode =
        readFile(
            "../shaders/bin/textured.frag.spv"
        );

    m_vertShaderModule =
        createShaderModule(vertShaderCode);
    m_fragShaderModule =
        createShaderModule(fragShaderCode);

    std::cout << "Shaders loaded" << std::endl;

    VkPipelineShaderStageCreateInfo vertStageInfo{};
    vertStageInfo.sType =
        VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    vertStageInfo.stage = VK_SHADER_STAGE_VERTEX_BIT;
    vertStageInfo.module = m_vertShaderModule;
    vertStageInfo.pName = "main";

    VkPipelineShaderStageCreateInfo fragStageInfo{};
    fragStageInfo.sType =
        VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    fragStageInfo.stage =
        VK_SHADER_STAGE_FRAGMENT_BIT;
    fragStageInfo.module = m_fragShaderModule;
    fragStageInfo.pName = "main";

    VkPipelineShaderStageCreateInfo shaderStages[] =
    {
        vertStageInfo,
        fragStageInfo
    };

    // Vertex input
    VkVertexInputBindingDescription
        bindingDescription{};
    bindingDescription.binding = 0;
    bindingDescription.stride = sizeof(Vertex);
    bindingDescription.inputRate =
        VK_VERTEX_INPUT_RATE_VERTEX;

    std::array<
        VkVertexInputAttributeDescription,
        2
    > attributeDescriptions{};

    attributeDescriptions[0].binding = 0;
    attributeDescriptions[0].location = 0;
    attributeDescriptions[0].format =
        VK_FORMAT_R32G32_SFLOAT;
    attributeDescriptions[0].offset =
        offsetof(Vertex, pos);

    attributeDescriptions[1].binding = 0;
    attributeDescriptions[1].location = 1;
    attributeDescriptions[1].format =
        VK_FORMAT_R32G32_SFLOAT;
    attributeDescriptions[1].offset =
        offsetof(Vertex, uv);

    VkPipelineVertexInputStateCreateInfo
        vertexInputInfo{};
    vertexInputInfo.sType =
        VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
    vertexInputInfo
        .vertexBindingDescriptionCount = 1;
    vertexInputInfo
        .pVertexBindingDescriptions =
        &bindingDescription;
    vertexInputInfo
        .vertexAttributeDescriptionCount =
        static_cast<uint32_t>(
            attributeDescriptions.size()
        );
    vertexInputInfo
        .pVertexAttributeDescriptions =
        attributeDescriptions.data();

    // Input assembly
    VkPipelineInputAssemblyStateCreateInfo
        inputAssembly{};
    inputAssembly.sType =
        VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
    inputAssembly.topology =
        VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;

    // Viewport state (dynamic, values don't matter)
    VkViewport dummyViewport{};
    dummyViewport.x = 0.0f;
    dummyViewport.y = 0.0f;
    dummyViewport.width = 1.0f;
    dummyViewport.height = 1.0f;
    dummyViewport.minDepth = 0.0f;
    dummyViewport.maxDepth = 1.0f;

    VkRect2D dummyScissor{};
    dummyScissor.offset = {0, 0};
    dummyScissor.extent = {1, 1};

    VkPipelineViewportStateCreateInfo
        viewportState{};
    viewportState.sType =
        VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
    viewportState.viewportCount = 1;
    viewportState.pViewports = &dummyViewport;
    viewportState.scissorCount = 1;
    viewportState.pScissors = &dummyScissor;

    // Rasterizer
    VkPipelineRasterizationStateCreateInfo
        rasterizer{};
    rasterizer.sType =
        VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
    rasterizer.polygonMode = VK_POLYGON_MODE_FILL;
    rasterizer.lineWidth = 1.0f;
    rasterizer.cullMode = VK_CULL_MODE_BACK_BIT;
    rasterizer.frontFace =
        VK_FRONT_FACE_CLOCKWISE;

    // Multisampling
    VkPipelineMultisampleStateCreateInfo
        multisampling{};
    multisampling.sType =
        VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
    multisampling.rasterizationSamples =
        VK_SAMPLE_COUNT_1_BIT;

    // Dynamic state
    VkDynamicState dynamicStates[] =
    {
        VK_DYNAMIC_STATE_VIEWPORT,
        VK_DYNAMIC_STATE_SCISSOR
    };

    VkPipelineDynamicStateCreateInfo
        dynamicState{};
    dynamicState.sType =
        VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
    dynamicState.dynamicStateCount = 2;
    dynamicState.pDynamicStates = dynamicStates;

    // Color blend
    VkPipelineColorBlendAttachmentState
        colorBlendAttachment{};
    colorBlendAttachment.colorWriteMask =
        VK_COLOR_COMPONENT_R_BIT |
        VK_COLOR_COMPONENT_G_BIT |
        VK_COLOR_COMPONENT_B_BIT |
        VK_COLOR_COMPONENT_A_BIT;

    VkPipelineColorBlendStateCreateInfo
        colorBlending{};
    colorBlending.sType =
        VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
    colorBlending.attachmentCount = 1;
    colorBlending.pAttachments =
        &colorBlendAttachment;

    // Descriptor set layout
    VkDescriptorSetLayoutBinding
        samplerLayoutBinding{};
    samplerLayoutBinding.binding = 0;
    samplerLayoutBinding.descriptorCount = 1;
    samplerLayoutBinding.descriptorType =
        VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    samplerLayoutBinding.pImmutableSamplers =
        nullptr;
    samplerLayoutBinding.stageFlags =
        VK_SHADER_STAGE_FRAGMENT_BIT;

    VkDescriptorSetLayoutCreateInfo
        layoutInfo{};
    layoutInfo.sType =
        VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
    layoutInfo.bindingCount = 1;
    layoutInfo.pBindings = &samplerLayoutBinding;

    if (vkCreateDescriptorSetLayout(
            m_device,
            &layoutInfo,
            nullptr,
            &m_descriptorSetLayout) != VK_SUCCESS)
    {
        throw std::runtime_error(
            "Failed to create descriptor set layout"
        );
    }

    VkPipelineLayoutCreateInfo
        pipelineLayoutInfo{};
    pipelineLayoutInfo.sType =
        VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
    pipelineLayoutInfo.setLayoutCount = 1;
    pipelineLayoutInfo.pSetLayouts =
        &m_descriptorSetLayout;

    if (vkCreatePipelineLayout(
            m_device,
            &pipelineLayoutInfo,
            nullptr,
            &m_pipelineLayout) != VK_SUCCESS)
    {
        throw std::runtime_error(
            "Failed to create pipeline layout"
        );
    }

    // Graphics pipeline
    VkGraphicsPipelineCreateInfo
        pipelineInfo{};
    pipelineInfo.sType =
        VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
    pipelineInfo.stageCount = 2;
    pipelineInfo.pStages = shaderStages;
    pipelineInfo.pVertexInputState =
        &vertexInputInfo;
    pipelineInfo.pInputAssemblyState =
        &inputAssembly;
    pipelineInfo.pViewportState =
        &viewportState;
    pipelineInfo.pRasterizationState =
        &rasterizer;
    pipelineInfo.pMultisampleState =
        &multisampling;
    pipelineInfo.pDynamicState = &dynamicState;
    pipelineInfo.pColorBlendState =
        &colorBlending;
    pipelineInfo.layout = m_pipelineLayout;
    pipelineInfo.renderPass = m_renderPass;
    pipelineInfo.subpass = 0;

    if (vkCreateGraphicsPipelines(
            m_device,
            VK_NULL_HANDLE,
            1,
            &pipelineInfo,
            nullptr,
            &m_graphicsPipeline) != VK_SUCCESS)
    {
        throw std::runtime_error(
            "Failed to create graphics pipeline"
        );
    }

    std::cout << "Graphics pipeline created"
              << std::endl;
}

RenderPipeline::~RenderPipeline()
{
    release();
}

void RenderPipeline::release()
{
    if (m_device == VK_NULL_HANDLE)
    {
        return;
    }

    vkDestroyPipeline(
        m_device,
        m_graphicsPipeline,
        nullptr
    );
    m_graphicsPipeline = VK_NULL_HANDLE;

    vkDestroyPipelineLayout(
        m_device,
        m_pipelineLayout,
        nullptr
    );
    m_pipelineLayout = VK_NULL_HANDLE;

    vkDestroyDescriptorSetLayout(
        m_device,
        m_descriptorSetLayout,
        nullptr
    );
    m_descriptorSetLayout = VK_NULL_HANDLE;

    vkDestroyShaderModule(
        m_device,
        m_vertShaderModule,
        nullptr
    );
    m_vertShaderModule = VK_NULL_HANDLE;

    vkDestroyShaderModule(
        m_device,
        m_fragShaderModule,
        nullptr
    );
    m_fragShaderModule = VK_NULL_HANDLE;

    vkDestroyRenderPass(
        m_device,
        m_renderPass,
        nullptr
    );
    m_renderPass = VK_NULL_HANDLE;

    m_device = VK_NULL_HANDLE;
}