#include <BlockAtlas/BlockAtlas.hpp>
#include <ChunkEncoder/ChunkEncoder.hpp>
#include <Chunk/Chunk.hpp>
#include <BlockRegistry/BlockRegistry.hpp>
#include <ChunkWindow/ChunkWindow.hpp>
#include <Display/WidgetDisplay/WorldDisplay/WorldDisplay.hpp>
#include <PipelineStage/LightingStage/LightingStage.hpp>
#include <Utilities/FixedPoint.hpp>
#include <Vk/ShaderConst.hpp>
#include <World/World.hpp>

Block& ChunkEncoder::getBlockById(int id)
{
    return window.world.getBlockById(id);
}

VoxelTree::NodeID ChunkEncoder::encodeNode(int level, glm::ivec3 offset)
{
    VoxelTree &svo = window.svo;
    if (level == 0)
    {
        return spriteGrid[offset.z * 16 * 16 + offset.y * 16 + offset.x];
    }

    uint32_t childNodes[8] = {};
    bool nodeNeeded = false;

    for (uint32_t childIndex = 0; childIndex < 8; childIndex++)
    {
        glm::ivec3 childOffset = offset << 1;
        childOffset.x |= childIndex & 1;
        childOffset.y |= (childIndex >> 1) & 1;
        childOffset.z |= (childIndex >> 2) & 1;

        childNodes[childIndex] = encodeNode(level - 1, childOffset);
        if (childNodes[childIndex] != childNodes[0] || (childNodes[childIndex] & GLSL_SVO_VOXEL_BIT) == 0)
        {
            nodeNeeded = true;
        }
    }

    if (!nodeNeeded)
    {
        return childNodes[0];
    }

    VoxelTree::NodeID nid = svo.createNode(0);
    for (uint32_t childIndex = 0; childIndex < 8; childIndex++)
    {
        svo.setChildNode(nid, childIndex, childNodes[childIndex]);
    }

    allocatedNodes.push_back(nid);
    return nid;
}

void ChunkEncoder::updateBlockSound(int offsetX, int offsetY, int offsetZ, Sound &sound)
{
    int blockIndex = 16 * 16 * offsetZ + 16 * offsetY + offsetX;

    if (blockAudioSources.count(blockIndex) == 0 && sound.isNone())
    {
        // Don't create an audio source if we won't be playing anything anyway.
        return;
    }

    std::shared_ptr<AudioSource> &audioSource = blockAudioSources[blockIndex];
    if (audioSource == nullptr)
    {
        // The audio source does not yet exist, create it.
        audioSource = std::make_shared<AudioSource>();
        audioSource->setPosition(glm::vec3(
            chunk.base.x + offsetX + 0.5f,
            chunk.base.y + offsetY + 0.5f,
            chunk.base.z + offsetZ + 0.5f
        ));
    }

    audioSource->setSound(sound);
}

void ChunkEncoder::updateBlockEffectTable(int offsetX, int offsetY, int offsetZ, Block &block, uint8_t meta)
{
    int blockIndex = 16 * 16 * offsetZ + 16 * offsetY + offsetX;

    std::vector<std::shared_ptr<BlockEffect>> effects;
    block.getEffects(meta, glm::vec3(offsetX, offsetY, offsetZ), effects);

    if (effects.size() == 0)
    {
        auto it = blockEffects.find(blockIndex);
        if (it != blockEffects.end())
        {
            blockEffects.erase(it);
        }
    }
    else
    {
        blockEffects[blockIndex] = effects;
    }
}

void ChunkEncoder::updateBlockLight(int offsetX, int offsetY, int offsetZ, Block &block, uint8_t meta)
{
    glm::vec3 lightColor = block.getLightColor(meta);
    float variance = block.getLightVariance(meta);
    int tableIndex = (offsetZ << 8) | (offsetY << 4) | offsetX;

    if (glm::length(lightColor) + variance > 0.001f)
    {
        if (blockLightIds.count(tableIndex) == 0)
        {
            blockLightIds[tableIndex] = window.addBlockLight(
                LightingStage::LightSource(
                    (glm::ivec3(offsetX, offsetY, offsetZ) + chunk.base) * 8 + block.getLightOffset(meta),
                    lightColor,
                    variance
                )
            );
        }
        else
        {
            window.setBlockLight(
                blockLightIds[tableIndex],
                LightingStage::LightSource(
                    (glm::ivec3(offsetX, offsetY, offsetZ) + chunk.base) * 8 + block.getLightOffset(meta),
                    lightColor,
                    variance
                )
            );
        }
    }
    else
    {
        auto currentLight = blockLightIds.find(tableIndex);
        if (currentLight != blockLightIds.end())
        {
            window.releaseBlockLight(currentLight->second);
            blockLightIds.erase(currentLight);
        }
    }
}

VoxelTree::NodeID ChunkEncoder::detachRecur(VoxelTree::Order order, VoxelTree::NodeID nid)
{
    VoxelTree &svo = window.svo;

    if (order == 7)
    {
        return GLSL_SVO_VOXEL_AIR;
    }

    if (nid & GLSL_SVO_VOXEL_BIT)
    {
        return GLSL_SVO_VOXEL_AIR;
    }

    uint32_t childIndex = VoxelTree::getChildNodeIndex(chunk.base << 3, order);
    VoxelTree::NodeID childNode = svo.getChildNode(nid, childIndex);

    childNode = detachRecur(order - 1, childNode);
    svo.setChildNode(nid, childIndex, childNode);

    if (nid == VoxelTree::rootNode)
    {
        // We cannot replace the root node.
        return nid;
    }

    return svo.releaseNodeIfUniform(nid);
}

ChunkEncoder::ChunkEncoder(ChunkWindow &window, Chunk &chunk) : window(window), chunk(chunk) {}

ChunkEncoder::~ChunkEncoder()
{
    if (blockLightIds.size() != 0)
    {
        crash("Sanity check failed: ChunkEncoder deleted with attached light sources.");
    }

    if (allocatedNodes.size() != 0)
    {
        crash("Sanity check failed: ChunkEncoder deleted with allocated SVO nodes.");
    }
}

void ChunkEncoder::generateGrid(const DenseArray &variantPalette, const DenseArray &blockGrid)
{
    blockAudioSources.clear();
    blockEffects.clear();
    initialLightSources.clear();

    for (size_t index = 0; index < 0x1000; index++)
    {
        glm::ivec3 offset(index & 0xF, (index >> 4) & 0xF, (index >> 8) & 0xF);
        uint32_t paletteIndex = blockGrid.get(index);

        uint32_t variant = paletteIndex;
        if (variantPalette.getNumEntries() != 0)
        {
            variant = variantPalette.get(paletteIndex);
        }

        int32_t blockId = variant >> 8;
        uint8_t meta = variant & 0xFF;

        Block &block = getBlockById(blockId);
        BlockAtlas::SpriteID sprite;
        sprite = block.getBlockSprite(offset + chunk.base, meta, window.world);

        spriteGrid[index] = sprite;
        updateBlockSound(offset.x, offset.y, offset.z, block.getBackgroundSound(meta));
        updateBlockEffectTable(offset.x, offset.y, offset.z, block, meta);

        glm::vec3 lightColor = block.getLightColor(meta);
        float variance = block.getLightVariance(meta);
        if (glm::length(lightColor) + variance > 0.001f)
        {
            initialLightSources.push_back(LightingStage::LightSource(
                (offset + chunk.base) * 8 + block.getLightOffset(meta),
                lightColor,
                variance
            ));
        }
    }
}

void ChunkEncoder::encode()
{
    for (LightingStage::LightSource &lightSource : initialLightSources)
    {
        glm::ivec3 blockPos = (lightSource.sourceVoxelAndRadius >> 3u) & 0xFu;
        int index = (blockPos.z << 8) | (blockPos.y << 4) | blockPos.x;
        blockLightIds[index] = window.addBlockLight(lightSource);
    }

    initialLightSources.clear();

    VoxelTree::NodeID chunkNode = encodeNode(4, glm::ivec3(0));

    VoxelTree::Order order = VoxelTree::rootOrder;
    VoxelTree::NodeID nid = VoxelTree::rootNode;
    VoxelTree &svo = window.svo;
    
    for (;;)
    {
        uint32_t childIndex = VoxelTree::getChildNodeIndex(chunk.base << 3, order);

        if (order == 8)
        {
            svo.setChildNode(nid, childIndex, chunkNode);
            break;
        }

        VoxelTree::NodeID childNode = svo.getChildNode(nid, childIndex);
        if (childNode == chunkNode)
        {
            // The chunk node is just one voxel, and this voxel is already inserted at this position.
            // We have nothing more to do.
            break;
        }

        if (childNode & GLSL_SVO_VOXEL_BIT)
        {
            childNode = svo.createNode(childNode);
            svo.setChildNode(nid, childIndex, childNode);
        }

        nid = childNode;
        order--;
    }
}

void ChunkEncoder::refreshBlock(glm::ivec3 pos)
{
    Block *block = window.world.getBlockAt(pos);
    if (block == nullptr)
    {
        return;
    }

    uint8_t meta = *window.world.getBlockMeta(pos);
    BlockAtlas::SpriteID sprite = block->getBlockSprite(pos, meta, window.world);

    updateBlockSound(pos.x & 0xF, pos.y & 0xF, pos.z & 0xF, block->getBackgroundSound(meta));
    updateBlockEffectTable(pos.x & 0xF, pos.y & 0xF, pos.z & 0xF, *block, meta);
    updateBlockLight(pos.x & 0xF, pos.y & 0xF, pos.z & 0xF, *block, meta);

    VoxelTree::Order order = VoxelTree::rootOrder;
    VoxelTree::NodeID nid = VoxelTree::rootNode;
    VoxelTree &svo = window.svo;
    
    for (;;)
    {
        uint32_t childIndex = VoxelTree::getChildNodeIndex(pos << 3, order);

        if (order == 4)
        {
            svo.setChildNode(nid, childIndex, sprite);
            break;
        }

        VoxelTree::NodeID childNode = svo.getChildNode(nid, childIndex);
        if (childNode & GLSL_SVO_VOXEL_BIT)
        {
            childNode = svo.createNode(childNode);
            svo.setChildNode(nid, childIndex, childNode);

            if (order <= 7)
            {
                allocatedNodes.push_back(childNode);
            }
        }

        nid = childNode;
        order--;
    }
}

float ChunkEncoder::getEncodingRelativeDeadline()
{
    glm::vec3 delta(wfvec3(chunk.base) - window.player->pos);
    return glm::length(delta) / (32.0f * 16.0);
}

void ChunkEncoder::detach()
{
    detachRecur(VoxelTree::rootOrder, VoxelTree::rootNode);

    for (auto &[_, lightId] : blockLightIds)
    {
        window.releaseBlockLight(lightId);
    }

    for (VoxelTree::NodeID nid : allocatedNodes)
    {
        window.svo.releaseNode(nid);
    }

    blockLightIds.clear();
    allocatedNodes.clear();
}

void ChunkEncoder::updateEffects(ChunkWindow &window)
{
    if (blockEffects.size() == 0)
    {
        return;
    }

    for (auto &[_, effects] : blockEffects)
    {
        for (auto &effect : effects)
        {
            effect->update(window);
        }
    }
}

bool ChunkEncoder::hasBlockEffects()
{
    return !blockEffects.empty();
}