#include <AssetManager/AssetManager.hpp>
#include <Cubemap/Cubemap.hpp>
#include <Log/Log.hpp>
#include <Utilities/Utilities.hpp>
#include <GraphicsContext/GraphicsContext.hpp>
#include <Vk/ShaderConst.hpp>
#include <vulkan/vulkan_core.h>

#define DECL_CUBEMAP(name) Cubemap Cubemap::name(#name)
#include <Cubemap/CubemapList.hpp>

Cubemap::Cubemap(const std::string &name) : m_name(name) {}

void Cubemap::load(std::atomic<unsigned int> &tasksDone)
{
    Log::info(formatString("Loading cubemap '$+' ...", m_name));
    std::shared_ptr<Vk::Device> device = GraphicsContext::instance->getDevice();

    static const char* faceNames[] = {
        "right.jpg",
        "left.jpg",
        "top.jpg",
        "bottom.jpg",
        "front.jpg",
        "back.jpg",
    };

    std::shared_ptr<Vk::Buffer> stagingBuffers[6];
    std::shared_ptr<Vk::CommandBuffer> cmdbuf = GraphicsContext::instance->createAsyncCommandBuffer();
    cmdbuf->reset();
    cmdbuf->begin();

    int i;
    for (i=0; i<6; i++)
    {
        std::string faceImagePath = assets.resolve(formatString("cubemaps/$+/$+", m_name, faceNames[i]));

        int comp, width, height;
        unsigned char *pixels = stbi_load(faceImagePath.c_str(), &width, &height, &comp, STBI_rgb_alpha);
        
        if (pixels == nullptr)
        {
            crash(formatString("Failed to load image file '$+'!", faceImagePath));
        }

        if (m_image == nullptr)
        {
            // If the image was not yet created we must create it now that we know the width and height.
            m_image = std::make_shared<Vk::Image>(
                device,
                VK_IMAGE_TYPE_2D,
                width,
                height,
                VK_FORMAT_R8G8B8A8_SRGB,
                VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT,
                VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
                6,
                VK_IMAGE_CREATE_CUBE_COMPATIBLE_BIT
            );

            m_imageView = std::make_shared<Vk::ImageView>(
                device,
                m_image,
                VK_IMAGE_VIEW_TYPE_CUBE,
                VK_IMAGE_ASPECT_COLOR_BIT
            );

            cmdbuf->transitionImageLayout(*m_image, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
        }

        stagingBuffers[i] = std::make_shared<Vk::Buffer>(
            device,
            4 * width * height,
            VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
        );

        std::memcpy(stagingBuffers[i]->hostdata(), pixels, 4 * width * height);
        stbi_image_free(pixels);

        cmdbuf->copyBufferToImage(*stagingBuffers[i], *m_image, width, height, i);
    }

    cmdbuf->transitionImageLayout(*m_image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
    cmdbuf->end();

    GraphicsContext::onNextFrame([&tasksDone, stagingBuffers, cmdbuf](GraphicsContext &context)
    {
        std::shared_ptr<Vk::Fence> fence = std::make_shared<Vk::Fence>(context.getDevice());
        cmdbuf->submit(fence.get());
        fence->wait();

        tasksDone++;
    });
}

void Cubemap::generate(glm::uvec2 size, Vk::ShaderModule &shader, std::atomic<unsigned int> &tasksDone)
{
    Log::info(formatString("Generating cubemap '$+' ...", m_name));
    auto device = GraphicsContext::instance->getDevice();

    auto cmdbuf = GraphicsContext::instance->createAsyncCommandBuffer();
    cmdbuf->reset();
    cmdbuf->begin();

    m_image = std::make_shared<Vk::Image>(
        device,
        VK_IMAGE_TYPE_2D,
        size.x,
        size.y,
        VK_FORMAT_R32G32B32A32_SFLOAT,
        VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_STORAGE_BIT,
        VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
        6,
        VK_IMAGE_CREATE_CUBE_COMPATIBLE_BIT
    );

    m_imageView = std::make_shared<Vk::ImageView>(
        device,
        m_image,
        VK_IMAGE_VIEW_TYPE_CUBE,
        VK_IMAGE_ASPECT_COLOR_BIT
    );

    cmdbuf->barrier(
        *m_image,
        VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
        0,
        VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
        VK_ACCESS_SHADER_WRITE_BIT,
        VK_IMAGE_LAYOUT_UNDEFINED,
        VK_IMAGE_LAYOUT_GENERAL,
        VK_IMAGE_ASPECT_COLOR_BIT
    );

    auto cubemapDescSetLayout = Vk::DescriptorSetLayout::Builder{}
        .addBinding(
            0,
            VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
            1,
            VK_SHADER_STAGE_COMPUTE_BIT
        )
        .build(GraphicsContext::instance->getDevice());
    
    auto cubemapDescPool = std::make_shared<Vk::DescriptorPool>(
        device,
        VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
        1,
        1
    );

    auto cubemapDescSet = std::make_shared<Vk::DescriptorSet>(
        device,
        cubemapDescPool,
        cubemapDescSetLayout
    );

    cubemapDescSet->setStorageImage(0, *m_imageView);

    auto cubemapGeneratePipeline = Vk::Pipeline::ComputeBuilder{}
        .addDescriptorSetLayout(cubemapDescSetLayout)
        .build(GraphicsContext::instance->getDevice(), shader);
    
    glm::uvec2 tiles = (size + glm::uvec2(GLSL_CUBEMAP_GEN_TILE_SIZE - 1)) / glm::uvec2(GLSL_CUBEMAP_GEN_TILE_SIZE);

    cmdbuf->bindPipeline(VK_PIPELINE_BIND_POINT_COMPUTE, *cubemapGeneratePipeline);
    cmdbuf->bindDescriptorSet(
        VK_PIPELINE_BIND_POINT_COMPUTE,
        *cubemapGeneratePipeline,
        0,
        *cubemapDescSet
    );
    cmdbuf->dispatch(glm::uvec3(tiles, 1));
    
    cmdbuf->barrier(
        *m_image,
        VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
        VK_ACCESS_SHADER_WRITE_BIT,
        VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
        VK_ACCESS_SHADER_READ_BIT,
        VK_IMAGE_LAYOUT_GENERAL,
        VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
        VK_IMAGE_ASPECT_COLOR_BIT
    );

    cmdbuf->end();

    GraphicsContext::onNextFrame([&tasksDone, cmdbuf, cubemapDescSetLayout, cubemapDescSet, cubemapDescPool, cubemapGeneratePipeline](GraphicsContext &context)
    {
        std::shared_ptr<Vk::Fence> fence = std::make_shared<Vk::Fence>(context.getDevice());
        cmdbuf->submit(fence.get());
        fence->wait();

        tasksDone++;
    });
}

Vk::ImageView& Cubemap::imageView()
{
    return *m_imageView;
}