#include <destrum/Graphics/Pipelines/LineRenderingPass.h>

#include <algorithm>
#include <cstdint>
#include <limits>
#include <stdexcept>

#include <destrum/FS/AssetFS.h>
#include <destrum/Graphics/Camera.h>
#include <destrum/Graphics/GPUImage.h>
#include <destrum/Graphics/GfxDevice.h>
#include <destrum/Graphics/Pipeline.h>
#include <destrum/Graphics/Util.h>

#include "volk.h"

namespace {
    constexpr std::size_t InitialVertexCapacity = 256;
}

void LineRenderingPass::init(GfxDevice& gfxDevice, VkFormat drawImageFormat)
{
    try {
        vertexBuffer.init(
            gfxDevice,
            VK_BUFFER_USAGE_VERTEX_BUFFER_BIT,
            InitialVertexCapacity * sizeof(LineVertex),
            "line vertices");
        vertexCapacity = InitialVertexCapacity;

        const auto vertexShader =
            AssetFS::GetInstance().GetCookedPathForFile("engine://shaders/line.vert");
        const auto fragmentShader =
            AssetFS::GetInstance().GetCookedPathForFile("engine://shaders/line.frag");

        const auto pushConstantRange = VkPushConstantRange{
            .stageFlags = VK_SHADER_STAGE_VERTEX_BIT,
            .offset = 0,
            .size = sizeof(glm::mat4),
        };

        const auto pipelineLayoutInfo = VkPipelineLayoutCreateInfo{
            .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
            .pushConstantRangeCount = 1,
            .pPushConstantRanges = &pushConstantRange,
        };

        VK_CHECK(vkCreatePipelineLayout(
            gfxDevice.getDevice(),
            &pipelineLayoutInfo,
            nullptr,
            &pipelineLayout));

        PipelineConfigInfo pipelineConfig{};
        Pipeline::DefaultPipelineConfigInfo(pipelineConfig);
        pipelineConfig.name = "Line Rendering Pipeline";
        pipelineConfig.pipelineLayout = pipelineLayout;
        pipelineConfig.inputAssemblyInfo.topology = VK_PRIMITIVE_TOPOLOGY_LINE_LIST;
        pipelineConfig.rasterizationInfo.cullMode = VK_CULL_MODE_NONE;
        pipelineConfig.depthStencilInfo.depthTestEnable = VK_FALSE;
        pipelineConfig.depthStencilInfo.depthWriteEnable = VK_FALSE;
        pipelineConfig.depthStencilInfo.depthCompareOp = VK_COMPARE_OP_ALWAYS;
        pipelineConfig.colorAttachments = {drawImageFormat};
        pipelineConfig.depthAttachment = VK_FORMAT_UNDEFINED;

        pipelineConfig.vertexBindingDescriptions = {
            VkVertexInputBindingDescription{
                .binding = 0,
                .stride = sizeof(LineVertex),
                .inputRate = VK_VERTEX_INPUT_RATE_VERTEX,
            }
        };
        pipelineConfig.vertexAttributeDescriptions = {
            VkVertexInputAttributeDescription{
                .location = 0,
                .binding = 0,
                .format = VK_FORMAT_R32G32B32A32_SFLOAT,
                .offset = offsetof(LineVertex, position),
            },
            VkVertexInputAttributeDescription{
                .location = 1,
                .binding = 0,
                .format = VK_FORMAT_R32G32B32A32_SFLOAT,
                .offset = offsetof(LineVertex, color),
            },
        };

        pipelineConfig.colorBlendAttachment.blendEnable = VK_TRUE;
        pipelineConfig.colorBlendAttachment.srcColorBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA;
        pipelineConfig.colorBlendAttachment.dstColorBlendFactor =
            VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA;
        pipelineConfig.colorBlendAttachment.colorBlendOp = VK_BLEND_OP_ADD;
        pipelineConfig.colorBlendAttachment.srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE;
        pipelineConfig.colorBlendAttachment.dstAlphaBlendFactor =
            VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA;
        pipelineConfig.colorBlendAttachment.alphaBlendOp = VK_BLEND_OP_ADD;

        pipeline = std::make_unique<Pipeline>(
            gfxDevice,
            vertexShader.string(),
            fragmentShader.string(),
            pipelineConfig);
    } catch (...) {
        cleanup(gfxDevice);
        throw;
    }
}

void LineRenderingPass::draw(
    VkCommandBuffer cmd,
    GfxDevice& gfxDevice,
    const GPUImage& target,
    const Camera& camera,
    LineRenderingManager& lineManager)
{
    const auto& lines = lineManager.GetLines();
    if (!pipeline || lines.empty()) {
        lineManager.ClearLines();
        return;
    }

    if (lines.size() > std::numeric_limits<std::size_t>::max() / 2) {
        throw std::overflow_error("Too many lines submitted for rendering");
    }

    std::vector<LineVertex> vertices;
    vertices.reserve(lines.size() * 2);
    for (const Line& line : lines) {
        vertices.push_back(LineVertex{
            .position = glm::vec4{line.start, 1.0f},
            .color = line.color,
        });
        vertices.push_back(LineVertex{
            .position = glm::vec4{line.end, 1.0f},
            .color = line.color,
        });
    }

    if (vertices.size() > std::numeric_limits<std::uint32_t>::max()) {
        throw std::overflow_error("Too many line vertices submitted for rendering");
    }

    ensureVertexCapacity(gfxDevice, vertices.size());
    vertexBuffer.uploadNewData(
        cmd,
        gfxDevice.getCurrentFrameIndex(),
        vertices.data(),
        vertices.size() * sizeof(LineVertex));

    auto renderInfo = vkutil::createRenderingInfo({
        .renderExtent = target.getExtent2D(),
        .colorImageView = target.imageView,
    });

    vkCmdBeginRendering(cmd, &renderInfo.renderingInfo);

    pipeline->bind(cmd);

    const auto viewport = VkViewport{
        .x = 0.0f,
        .y = 0.0f,
        .width = static_cast<float>(target.extent.width),
        .height = static_cast<float>(target.extent.height),
        .minDepth = 0.0f,
        .maxDepth = 1.0f,
    };
    vkCmdSetViewport(cmd, 0, 1, &viewport);

    const auto scissor = VkRect2D{
        .offset = {},
        .extent = target.getExtent2D(),
    };
    vkCmdSetScissor(cmd, 0, 1, &scissor);
    vkCmdSetPolygonModeEXT(cmd, VK_POLYGON_MODE_FILL);

    const glm::mat4 viewProjection = camera.GetViewProjectionMatrix();
    vkCmdPushConstants(
        cmd,
        pipelineLayout,
        VK_SHADER_STAGE_VERTEX_BIT,
        0,
        sizeof(viewProjection),
        &viewProjection);

    const VkBuffer vertexBufferHandle = vertexBuffer.getBuffer().buffer;
    const VkDeviceSize vertexBufferOffset = 0;
    vkCmdBindVertexBuffers(cmd, 0, 1, &vertexBufferHandle, &vertexBufferOffset);
    vkCmdDraw(cmd, static_cast<std::uint32_t>(vertices.size()), 1, 0, 0);

    vkCmdEndRendering(cmd);
    lineManager.ClearLines();
}

void LineRenderingPass::cleanup(GfxDevice& gfxDevice)
{
    pipeline.reset();

    if (pipelineLayout != VK_NULL_HANDLE && gfxDevice.getDevice() != VK_NULL_HANDLE) {
        vkDestroyPipelineLayout(gfxDevice.getDevice(), pipelineLayout, nullptr);
    }
    pipelineLayout = VK_NULL_HANDLE;

    vertexBuffer.cleanup(gfxDevice);
    vertexCapacity = 0;
}

void LineRenderingPass::ensureVertexCapacity(
    GfxDevice& gfxDevice,
    std::size_t requiredVertices)
{
    if (requiredVertices <= vertexCapacity) {
        return;
    }

    std::size_t newCapacity = std::max(vertexCapacity, InitialVertexCapacity);
    while (newCapacity < requiredVertices) {
        if (newCapacity > std::numeric_limits<std::size_t>::max() / 2) {
            newCapacity = requiredVertices;
            break;
        }
        newCapacity *= 2;
    }

    vertexBuffer.cleanup(gfxDevice);
    vertexBuffer.init(
        gfxDevice,
        VK_BUFFER_USAGE_VERTEX_BUFFER_BIT,
        newCapacity * sizeof(LineVertex),
        "line vertices");
    vertexCapacity = newCapacity;
}