#include <destrum/Graphics/Resources/Cubemap.h>

#include <destrum/Graphics/GfxDevice.h>
#include <destrum/Graphics/RenderResources.h>
#include <destrum/Graphics/Pipeline.h>
#include <destrum/Graphics/Util.h>

#include <glm/ext/matrix_clip_space.hpp>

#include <cmath>
#include <stdexcept>
#include <utility>

#include "volk.h"
#include "spdlog/spdlog.h"

CubeMap::CubeMap()
{
    m_projection = glm::perspective(
        glm::radians(90.0f),
        1.0f,
        0.1f,
        10.0f
    );
}

CubeMap::~CubeMap() = default;

void CubeMap::cleanup(GfxDevice& gfxDevice)
{
    VkDevice device = gfxDevice.getDevice();

    m_cubemapPipeline.reset();

    if (m_cubemapPipelineLayout != VK_NULL_HANDLE) {
        vkDestroyPipelineLayout(device, m_cubemapPipelineLayout, nullptr);
        m_cubemapPipelineLayout = VK_NULL_HANDLE;
    }

    m_cubemapImageID = NULL_IMAGE_ID;
    m_hdrImage = NULL_IMAGE_ID;
}

void CubeMap::LoadCubeMap(
    GfxDevice& gfxDevice,
    RenderResources& resources,
    const std::filesystem::path& directoryPath)
{
    m_hdrImage = resources.loadImageFromFile(
        gfxDevice,
        directoryPath,
        VK_IMAGE_USAGE_SAMPLED_BIT |
        VK_IMAGE_USAGE_TRANSFER_DST_BIT |
        VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,
        false
    );
}

void CubeMap::RenderToCubemap(
    GfxDevice& gfxDevice,
    RenderResources& resources,
    ImageID inputImage,
    VkImage outputImage,
    std::array<VkImageView, 6> faceViews,
    std::uint32_t size)
{
    if (!m_cubemapPipeline || m_cubemapPipelineLayout == VK_NULL_HANDLE) {
        throw std::runtime_error(
            "Cubemap pipeline not initialized. Call InitCubemapPipeline first."
        );
    }

    gfxDevice.getImmediateExecuter().immediateSubmit([&](VkCommandBuffer cmd) {
        vkutil::ScopedCmdLabel label(cmd, "HDRI To Cubemap", vkutil::labelColors::sky);

        VkImageMemoryBarrier barrier{VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER};
        barrier.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED;
        barrier.newLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
        barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        barrier.image = outputImage;
        barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        barrier.subresourceRange.baseMipLevel = 0;
        barrier.subresourceRange.levelCount = 1;
        barrier.subresourceRange.baseArrayLayer = 0;
        barrier.subresourceRange.layerCount = 6;
        barrier.srcAccessMask = 0;
        barrier.dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;

        vkCmdPipelineBarrier(
            cmd,
            VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
            VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
            0,
            0, nullptr,
            0, nullptr,
            1, &barrier
        );

        VkViewport viewport{};
        viewport.x = 0.0f;
        viewport.y = 0.0f;
        viewport.width = static_cast<float>(size);
        viewport.height = static_cast<float>(size);
        viewport.minDepth = 0.0f;
        viewport.maxDepth = 1.0f;

        VkRect2D scissor{};
        scissor.offset = {0, 0};
        scissor.extent = {size, size};

        for (std::uint32_t face = 0; face < 6; ++face) {
            vkutil::cmdBeginLabelFormatted(cmd, vkutil::labelColors::draw, "Cubemap Face %u", face);

            PC pc{};
            pc.viewMtx = viewMatrices[face];
            pc.projMtx = m_projection;
            pc.inputImageId = static_cast<std::uint32_t>(inputImage);

            VkRenderingAttachmentInfoKHR colorAttachment{
                VK_STRUCTURE_TYPE_RENDERING_ATTACHMENT_INFO_KHR
            };
            colorAttachment.imageView = faceViews[face];
            colorAttachment.imageLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
            colorAttachment.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
            colorAttachment.storeOp = VK_ATTACHMENT_STORE_OP_STORE;
            colorAttachment.clearValue.color = {0.f, 0.f, 0.f, 1.f};

            VkRenderingInfoKHR renderingInfo{
                VK_STRUCTURE_TYPE_RENDERING_INFO_KHR
            };
            renderingInfo.renderArea.offset = {0, 0};
            renderingInfo.renderArea.extent = {size, size};
            renderingInfo.layerCount = 1;
            renderingInfo.colorAttachmentCount = 1;
            renderingInfo.pColorAttachments = &colorAttachment;

            vkCmdBeginRendering(cmd, &renderingInfo);

            vkCmdSetViewport(cmd, 0, 1, &viewport);
            vkCmdSetScissor(cmd, 0, 1, &scissor);

            m_cubemapPipeline->bind(cmd);
            vkCmdSetPolygonModeEXT(cmd, VK_POLYGON_MODE_FILL);

            resources.bindBindlessDescSet(cmd, m_cubemapPipelineLayout);

            vkCmdPushConstants(
                cmd,
                m_cubemapPipelineLayout,
                VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT,
                0,
                sizeof(PC),
                &pc
            );

            vkCmdDraw(cmd, 36, 1, 0, 0);

            vkCmdEndRendering(cmd);

            vkutil::cmdEndLabel(cmd);
        }

        barrier.oldLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
        barrier.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        barrier.srcAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
        barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;

        vkCmdPipelineBarrier(
            cmd,
            VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
            VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
            0,
            0, nullptr,
            0, nullptr,
            1, &barrier
        );
    });
}

void CubeMap::CreateCubeMap(
    GfxDevice& gfxDevice,
    RenderResources& resources)
{
    if (m_hdrImage == NULL_IMAGE_ID) {
        throw std::runtime_error(
            "Cannot create cubemap before loading HDR image. Call LoadCubeMap first."
        );
    }

    GPUImage cubeMapImage = gfxDevice.createImageRaw({
        .format = VK_FORMAT_R32G32B32A32_SFLOAT,
        .usage =
            VK_IMAGE_USAGE_SAMPLED_BIT |
            VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT |
            VK_IMAGE_USAGE_TRANSFER_DST_BIT |
            VK_IMAGE_USAGE_TRANSFER_SRC_BIT,
        .flags = VK_IMAGE_CREATE_CUBE_COMPATIBLE_BIT,
        .extent = VkExtent3D{
            .width = m_cubeMapSize,
            .height = m_cubeMapSize,
            .depth = 1,
        },
        .numLayers = 6,
        .mipMap = false,
        .isCubemap = true
    });
    cubeMapImage.debugName = "skybox cubemap";
    vkutil::addDebugLabel(gfxDevice.getDevice(), cubeMapImage.image, cubeMapImage.debugName.c_str());

    std::array<VkImageView, 6> faceViews{};

    try {
        for (std::uint32_t face = 0; face < 6; ++face) {
            VkImageViewCreateInfo viewInfo{};
            viewInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
            viewInfo.image = cubeMapImage.image;
            viewInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
            viewInfo.format = VK_FORMAT_R32G32B32A32_SFLOAT;
            viewInfo.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            viewInfo.subresourceRange.baseMipLevel = 0;
            viewInfo.subresourceRange.levelCount = 1;
            viewInfo.subresourceRange.baseArrayLayer = face;
            viewInfo.subresourceRange.layerCount = 1;

            VK_CHECK(vkCreateImageView(
                gfxDevice.getDevice(),
                &viewInfo,
                nullptr,
                &faceViews[face]));

            const std::string faceViewName = "skybox cubemap face " + std::to_string(face) + " (view)";
            vkutil::addDebugLabel(gfxDevice.getDevice(), faceViews[face], faceViewName.c_str());
        }

        spdlog::info("HDRI image id = {}", m_hdrImage);

        RenderToCubemap(
            gfxDevice,
            resources,
            m_hdrImage,
            cubeMapImage.image,
            faceViews,
            m_cubeMapSize
        );

        for (std::uint32_t face = 0; face < 6; ++face) {
            cubeMapImage.setLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, face);
        }

        m_cubemapImageID = resources.addImageToCache(std::move(cubeMapImage));

        for (VkImageView& view : faceViews) {
            if (view != VK_NULL_HANDLE) {
                vkDestroyImageView(gfxDevice.getDevice(), view, nullptr);
                view = VK_NULL_HANDLE;
            }
        }
    } catch (...) {
        for (VkImageView& view : faceViews) {
            if (view != VK_NULL_HANDLE) {
                vkDestroyImageView(gfxDevice.getDevice(), view, nullptr);
            }
        }
        gfxDevice.destroyImage(cubeMapImage);
        throw;
    }
}

ImageID CubeMap::GetCubeMapImageID() const
{
    return m_cubemapImageID;
}

void CubeMap::InitCubemapPipeline(
    GfxDevice& gfxDevice,
    RenderResources& resources,
    const std::string& vertPath,
    const std::string& fragPath)
{
    if (m_cubemapPipeline) {
        return;
    }

    VkDevice device = gfxDevice.getDevice();

    m_cubemapVert = vertPath;
    m_cubemapFrag = fragPath;

    VkPushConstantRange pushConstantRange{};
    pushConstantRange.stageFlags =
        VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT;
    pushConstantRange.offset = 0;
    pushConstantRange.size = sizeof(PC);

    const auto layouts = std::array{
        resources.getBindlessDescSetLayout()
    };

    VkPipelineLayoutCreateInfo pipelineLayoutInfo{
        VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO
    };
    pipelineLayoutInfo.setLayoutCount = static_cast<std::uint32_t>(layouts.size());
    pipelineLayoutInfo.pSetLayouts = layouts.data();
    pipelineLayoutInfo.pushConstantRangeCount = 1;
    pipelineLayoutInfo.pPushConstantRanges = &pushConstantRange;

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

    PipelineConfigInfo pipelineConfig{};
    Pipeline::DefaultPipelineConfigInfo(pipelineConfig);
    pipelineConfig.name = "HDRI To Cubemap Pipeline";

    pipelineConfig.vertexAttributeDescriptions = {};
    pipelineConfig.vertexBindingDescriptions = {};
    pipelineConfig.pipelineLayout = m_cubemapPipelineLayout;
    pipelineConfig.colorAttachments = {VK_FORMAT_R32G32B32A32_SFLOAT};
    pipelineConfig.depthAttachment = VK_FORMAT_UNDEFINED;
    pipelineConfig.depthStencilInfo.depthTestEnable = VK_FALSE;
    pipelineConfig.depthStencilInfo.depthWriteEnable = VK_FALSE;
    pipelineConfig.rasterizationInfo.cullMode = VK_CULL_MODE_NONE;

    m_cubemapPipeline = std::make_unique<Pipeline>(
        gfxDevice,
        vertPath,
        fragPath,
        pipelineConfig
    );
}