#include <PipelineStage/DrawStage/ModelStage/ModelStage.hpp>
#include <Vk/DescriptorPool/DescriptorPool.hpp>
#include <VoxelSprite/VoxelSprite.hpp>
#include <iostream>

#include <VoxelSpriteRenderer/VoxelSpriteRenderer.hpp>
#include <Utilities/Utilities.hpp>
#include <GraphicsContext/GraphicsContext.hpp>
#include <Utilities/Extern.hpp>
#include <FrameDrawer/FrameDrawer.hpp>
#include <Display/WidgetDisplay/WorldDisplay/WorldDisplay.hpp>
#include <memory>

VoxelSpriteRenderer::VoxelSpriteRenderer(VoxelSprite &sprite)
{
    std::shared_ptr<Vk::Device> device = GraphicsContext::instance->getDevice();
    bufferSize = sizeof(ModelStage::ModelBuffer) + sizeof(uint32_t) * sprite.getVolume();

    stagingBuffer = std::make_unique<Vk::Buffer>(
        device,
        bufferSize,
        VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
    );

    deviceLocalBuffer = std::make_unique<Vk::Buffer>(
        device,
        bufferSize,
        VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
        VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT
    );

    ModelStage::ModelBuffer *modelBuffer = reinterpret_cast<ModelStage::ModelBuffer*>(stagingBuffer->hostdata());
    glm::ivec3 size = sprite.getSize();
    modelBuffer->xsize = size.x;
    modelBuffer->ysize = size.y;
    modelBuffer->zsize = size.z;

    glm::ivec3 offset;
    for (offset.z = 0; offset.z < size.z; offset.z++)
    {
        for (offset.y = 0; offset.y < size.y; offset.y++)
        {
            for (offset.x = 0; offset.x < size.x; offset.x++)
            {
                VoxelSprite::Color &color = sprite.at(offset);
                modelBuffer->voxels[offset.z * size.y * size.x + offset.y * size.x + offset.x] =
                    ((uint32_t) color.red << 24) |
                    ((uint32_t) color.green << 16) |
                    ((uint32_t) color.blue << 8) |
                    color.alpha;
            }
        }
    }
}

std::shared_ptr<Model> VoxelSpriteRenderer::makeModel()
{
    std::shared_ptr<Vk::Device> device = GraphicsContext::instance->getDevice();
    std::shared_ptr<Vk::DescriptorPool> descPool = Vk::DescriptorPool::Builder()
        .addDescriptorType(VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1)
        .build(device, 3);

    std::shared_ptr<Vk::CommandBuffer> cmdbuf = GraphicsContext::instance->createAsyncCommandBuffer();
    cmdbuf->reset();
    cmdbuf->begin();

    cmdbuf->transitionBufferToTransferDst(*deviceLocalBuffer);
    cmdbuf->copyBuffer(*stagingBuffer, 0, *deviceLocalBuffer, 0, bufferSize);
    cmdbuf->transitionBufferToShaderRead(*deviceLocalBuffer);
    cmdbuf->end();

    std::shared_ptr<Vk::Fence> fence = std::make_shared<Vk::Fence>(GraphicsContext::instance->getDevice());
    cmdbuf->submit(fence.get());
    fence->wait();

    descSet = std::make_shared<Vk::DescriptorSet>(
        device,
        descPool,
        WorldDisplay::modelDescriptorSetLayout
    );

    descSet->setStorageBuffer(0, *deviceLocalBuffer, 0, VK_WHOLE_SIZE);

    return std::make_shared<Model>(deviceLocalBuffer, descSet);
}