#include <BlockTintManager/BlockTintManager.hpp>
#include <FrameDrawer/FrameDrawer.hpp>
#include <GraphicsContext/GraphicsContext.hpp>
#include <Utilities/JHash.hpp>
#include <Display/WidgetDisplay/WorldDisplay/WorldDisplay.hpp>

uint32_t BlockTintManager::posHash(const glm::ivec3 &blockPos, uint32_t depth)
{
    uint32_t a = jhash(uint32_t(blockPos.x) * 73856093u + depth);
    uint32_t b = jhash(uint32_t(blockPos.y) * 19349663u + depth);
    uint32_t c = jhash(uint32_t(blockPos.z) * 83492791u + depth);

    return a ^ b ^ c;
}

uint32_t BlockTintManager::createHashNode()
{
    uint32_t index = uint32_t(hashTableNodes.size());
    hashTableNodes.push({});
    return index;
}

BlockTintManager::BlockTintManager(FrameDrawer &frameDrawer) :
    frameDrawer(frameDrawer),
    hashTableNodes(VK_BUFFER_USAGE_STORAGE_BUFFER_BIT),
    tintTable(VK_BUFFER_USAGE_STORAGE_BUFFER_BIT)
{
    hashTableNodes.push({});

    tintDescriptorPool = Vk::DescriptorPool::Builder()
        .addDescriptorType(VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 32)
        .build(GraphicsContext::instance->getDevice(), 16);
}

void BlockTintManager::setTintAt(const glm::ivec3 &blockPos, const glm::vec4 &tint)
{
    auto normBlockPos = glm::uvec3(blockPos) & glm::uvec3((1 << (GLSL_SVO_MAX_ORDER - 2)) - 1);
    uint32_t currentNodeIndex = 0;
    for (uint32_t depth = 0;; depth++)
    {
        HashTableNode &node = hashTableNodes[currentNodeIndex];
        uint32_t hash = posHash(normBlockPos, depth);

        uint32_t entryIndex = hash % GLSL_TINT_TABLE_NODE_SIZE;
        uint32_t entry = node.entries[entryIndex];

        if (entry & GLSL_TINT_TABLE_NODE_RESOLVED)
        {
            // If the tint entry has the same position, we update it. Otherwise, we need to create
            // a nested hash table to resolved the collision.
            uint32_t tintIndex = entry & ~GLSL_TINT_TABLE_NODE_RESOLVED;
            TintEntry &tintEntry = tintTable[tintIndex];

            if (glm::uvec3(tintEntry.blockPosAndTint) == normBlockPos)
            {
                tintEntry.blockPosAndTint.w = glm::packUnorm4x8(tint);
                return;
            }

            // Create a new hash table node to resolve the collision.
            uint32_t newNodeIndex = createHashNode();
            HashTableNode &newNode = hashTableNodes[newNodeIndex];
            newNode.entries[posHash(glm::ivec3(tintEntry.blockPosAndTint), depth + 1) % GLSL_TINT_TABLE_NODE_SIZE] = entry;
            node.entries[entryIndex] = newNodeIndex;

            currentNodeIndex = newNodeIndex;
        }
        else if (entry == 0)
        {
            uint32_t tintIndex = tintTable.size();
            tintTable.push({glm::uvec4(normBlockPos, glm::packUnorm4x8(tint))});
            node.entries[entryIndex] = tintIndex | GLSL_TINT_TABLE_NODE_RESOLVED;
            return;
        }
        else
        {
            currentNodeIndex = entry;
        }
    }
}

std::shared_ptr<Vk::DescriptorSet> BlockTintManager::getDescSet()
{
    auto descSet = std::make_shared<Vk::DescriptorSet>(
        GraphicsContext::instance->getDevice(),
        tintDescriptorPool,
        frameDrawer.worldDisplay.tintDescriptorSetLayout
    );

    auto hashTableBuffer = hashTableNodes.unpack();
    auto tintTableBuffer = tintTable.unpack();

    descSet->setStorageBuffer(GLSL_TINT_HASH_TABLE_BINDING_INDEX, *hashTableBuffer, 0, VK_WHOLE_SIZE);
    descSet->setStorageBuffer(GLSL_TINT_TABLE_BINDING_INDEX, *tintTableBuffer, 0, VK_WHOLE_SIZE);

    GraphicsContext::instance->holdResource(hashTableBuffer);
    GraphicsContext::instance->holdResource(tintTableBuffer);
    GraphicsContext::instance->holdResource(descSet);

    hashTableNodes.push({});

    return descSet;
}