#include <Utilities/JHash.hpp>
#include <VoxelTree/VoxelTree.hpp>
#include <cstdlib>
#include <iostream>
#include <BlockAtlas/BlockAtlas.hpp>
#include <Log/Log.hpp>
#include <Utilities/Utilities.hpp>
#include <Utilities/Extern.hpp>
#include <GraphicsContext/GraphicsContext.hpp>
#include <Vk/ShaderConst.hpp>
#include <Display/WidgetDisplay/WorldDisplay/WorldDisplay.hpp>
#include <BlockRegistry/BlockRegistry.hpp>
#include <vulkan/vulkan_core.h>

std::optional<VoxelTree::NodeID> BlockAtlas::tryCollapseNode(const VoxelTree::Node &node)
{
    for (size_t i = 1; i < 8; i++)
    {
        if (node.childRefs[i] != node.childRefs[0])
        {
            return std::nullopt;
        }
    }

    return node.childRefs[0];
}

BlockAtlas::NodeSample BlockAtlas::getNodeSample(VoxelTree::NodeID nodeId)
{
    if (nodeId & GLSL_SVO_VOXEL_BIT)
    {
        NodeSample sample;

        size_t numFrames = 1 << GLSL_SVO_VOXEL_FRAMEPOW_GET(nodeId);
        size_t voxelBaseIndex = nodeId & GLSL_SVO_VOXEL_INDEX_MASK;

        for (size_t frameno = 0; frameno < 8; frameno++)
        {
            size_t voxelIndex = voxelBaseIndex + (frameno & (numFrames-1));
            uint32_t radiance = voxelPalette[voxelIndex].lightSpec;
            RGBABytes filter = decodeColor(voxelPalette[voxelIndex].color);

            for (size_t channel = 0; channel < 3; channel++)
            {
                filter.channels[channel] = ((255 - filter.channels[channel]) * filter.alpha / 3) >> 8;
            }

            if (filter.alpha == 0xFF)
            {
                // Opaque; no light comes through.
                filter.asInt = 0xFFFFFF00; 
            }

            filter.alpha = 0;

            for (size_t face = 0; face < 6; face++)
            {
                sample.frames[frameno].faces[face].radiance = radiance;
                sample.frames[frameno].faces[face].filter = encodeColor(filter);
            }
        }

        return sample;
    }

    return nodeSamples[nodeId];
}

VoxelTree::NodeID BlockAtlas::addNode(const VoxelTree::Node &node)
{
    std::optional<VoxelTree::NodeID> collapsed = tryCollapseNode(node);
    if (collapsed.has_value())
    {
        return *collapsed;
    }

    uint32_t hash = 0;
    for (uint32_t ref : node.childRefs)
    {
        hash += ref;
    }

    size_t bucketIndex = hash % numNodeBuckets;
    for (SpriteID id : nodeHashSet[bucketIndex])
    {
        if (std::memcmp(&atlasNodes[id], &node, sizeof(VoxelTree::Node)) == 0)
        {
            return id;
        }
    }

    NodeSample nextOrderSample;

    NodeSample childSamples[8];
    for (size_t i = 0; i < 8; i++)
    {
        childSamples[i] = getNodeSample(node.childRefs[i]);
    }

    for (size_t frameno = 0; frameno < 8; frameno++)
    {
        for (size_t face = 0; face < 6; face++)
        {
            uint32_t dir = face & 1;
            uint32_t axis = (face >> 1) & 0x3;

            NodeSampleFace &combinedFace = nextOrderSample.frames[frameno].faces[face];
            glm::uvec3 combinedRadiance{};
            glm::uvec3 combinedFilter{};

            for (uint32_t childIndex = 0; childIndex < 8; childIndex++)
            {
                if (((childIndex >> axis) & 1) == dir)
                {
                    RGBABytes childRadiance = decodeColor(childSamples[childIndex].frames[frameno].faces[face].radiance);
                    for (size_t channel = 0; channel < 3; channel++)
                    {
                        combinedRadiance[channel] += childRadiance.channels[channel];
                    }
                }

                RGBABytes childFilter = decodeColor(childSamples[childIndex].frames[frameno].faces[face].filter);
                for (size_t channel = 0; channel < 3; channel++)
                {
                    combinedFilter[channel] += childFilter.channels[channel];
                }
            }

            RGBABytes normRadiance{};
            for (size_t i = 0; i < 3; i++)
            {
                normRadiance.channels[i] = combinedRadiance[i] >> 2;
            }

            RGBABytes normFilter{};
            for (size_t i = 0; i < 3; i++)
            {
                normFilter.channels[i] = combinedFilter[i] >> 3;
            }

            combinedFace.radiance = encodeColor(normRadiance);
            combinedFace.filter = encodeColor(normFilter);
        }
    }

    nodeSamples.push_back(nextOrderSample);
    atlasNodes.push_back(node);
    SpriteID id = atlasNodes.size() - 1;
    nodeHashSet[bucketIndex].push_back(id);
    return id;
}

RGBABytes BlockAtlas::decodeColor(uint32_t rgba)
{
    RGBABytes result;
    result.alpha = rgba;
    result.blue = rgba >> 8;
    result.green = rgba >> 16;
    result.red = rgba >> 24;
    return result;
}

uint32_t BlockAtlas::encodeColor(VoxelSprite::Color &color)
{
    return (((uint32_t) color.red) << 24) |
        (((uint32_t) color.green) << 16) |
        (((uint32_t) color.blue) << 8) |
        ((uint32_t) color.alpha);
}

uint32_t BlockAtlas::encodeVoxelFrames(const std::vector<Voxel> &voxelFrames)
{
    uint32_t hash = jhash(voxelFrames.size());
    for (const auto &voxel : voxelFrames)
    {
        hash ^= jhash(voxel.color);
        hash ^= jhash(voxel.bias);
        hash ^= jhash(voxel.lightSpec);
        hash ^= jhash(voxel.refract);
    }

    std::vector<uint32_t> &bucket = voxelHashSet[hash % numVoxelBuckets];
    for (uint32_t nodeRef : bucket)
    {
        if (areVoxelsSame(nodeRef, voxelFrames))
        {
            return nodeRef;
        }
    }

    int framePow;
    for (framePow = 0; (1 << framePow) != voxelFrames.size(); framePow++);

    size_t assignedIndex = voxelPalette.size();
    voxelPalette.insert(voxelPalette.end(), voxelFrames.begin(), voxelFrames.end());

    uint32_t nodeRef = GLSL_SVO_VOXEL_BIT | GLSL_SVO_VOXEL_FRAMEPOW_MAKE(framePow) | assignedIndex;
    bucket.push_back(nodeRef);

    return nodeRef;
}

VoxelTree::NodeID BlockAtlas::makeVoxelNode(BlockSprite &sprite, glm::uvec3 pos)
{
    std::vector<Voxel> renderedFrames;
    for (auto &frame : sprite.frames)
    {
        renderedFrames.push_back({});
        Voxel &voxel = renderedFrames.back();

        voxel.color = encodeColor(frame.colorMap.at(pos.x, pos.y, pos.z));
        voxel.frameDuration = sprite.frameDuration;

        if ((voxel.color & 0xFF) == 0)
        {
            voxel.color = 0;
            continue;
        }

        voxel.lightSpec = encodeColor(frame.lightSpecMap.at(pos.x, pos.y, pos.z));
        voxel.bias = encodeColor(frame.biasMap.at(pos.x, pos.y, pos.z));
        voxel.refract = encodeColor(frame.refractMap.at(pos.x, pos.y, pos.z));
    }

    // Collapse repeating animations into fewer frames.
    while (renderedFrames.size() != 1)
    {
        size_t half = renderedFrames.size() / 2;
        size_t i;

        for (i = 0; i < half; i++)
        {
            const Voxel &a = renderedFrames[i];
            const Voxel &b = renderedFrames[i + half];

            if (a.color != b.color || a.bias != b.bias || a.lightSpec != b.lightSpec || a.refract != b.refract)
            {
                break;
            }
        }

        if (i == half)
        {
            renderedFrames.resize(half);
        }
        else
        {
            break;
        }
    }

    if (renderedFrames.size() == 1 && (renderedFrames[0].color & 0xFF) == 0)
    {
        renderedFrames[0].frameDuration = BlockSprite::defaultFrameDuration;
    }

    return encodeVoxelFrames(renderedFrames);
}

VoxelTree::NodeID BlockAtlas::makeSpriteNode(BlockSprite &sprite, VoxelTree::Order order, glm::uvec3 offset)
{
    if (order == 0)
    {
        return makeVoxelNode(sprite, offset);
    }

    VoxelTree::Node node{};
    for (uint32_t childIndex = 0; childIndex < 8; childIndex++)
    {
        glm::uvec3 childOffset = offset << 1u;
        childOffset.x |= childIndex & 1;
        childOffset.y |= (childIndex >> 1) & 1;
        childOffset.z |= (childIndex >> 2) & 1;

        node.childRefs[childIndex] = makeSpriteNode(sprite, order - 1, childOffset);
    }

    return addNode(node);
}

bool BlockAtlas::areVoxelsSame(uint32_t nodeRef, const std::vector<Voxel> &voxels)
{
    int framePow;
    for (framePow = 0; (1 << framePow) != voxels.size(); framePow++);

    if (GLSL_SVO_VOXEL_FRAMEPOW_GET(nodeRef) != framePow)
    {
        return false;
    }

    size_t index = nodeRef & GLSL_SVO_VOXEL_INDEX_MASK;
    for (size_t offset = 0; offset < voxels.size(); offset++)
    {
        const Voxel &a = voxelPalette[index + offset];
        const Voxel &b = voxels[offset];

        if (a.color != b.color || a.bias != b.bias || a.lightSpec != b.lightSpec || a.refract != b.refract ||
            a.frameDuration != b.frameDuration)
        {
            return false;
        }
    }

    return true;
}

BlockAtlas::BlockAtlas()
{
    // The root node is node 0.
    atlasNodes.push_back({
        GLSL_SVO_VOXEL_AIR,
        GLSL_SVO_VOXEL_AIR,
        GLSL_SVO_VOXEL_AIR,
        GLSL_SVO_VOXEL_AIR,
        GLSL_SVO_VOXEL_AIR,
        GLSL_SVO_VOXEL_AIR,
        GLSL_SVO_VOXEL_AIR,
        GLSL_SVO_VOXEL_AIR,
    });

    // The sample for a root node will never be read.
    nodeSamples.push_back({});

    nodeHashSet.resize(numNodeBuckets);
    nodeHashSet[0].push_back(0);

    voxelHashSet.resize(numVoxelBuckets);

    std::vector<Voxel> airVoxel{{}};
    airVoxel[0].frameDuration = BlockSprite::defaultFrameDuration;
    uint32_t airNodeRef = encodeVoxelFrames(airVoxel);
    
    if (airNodeRef != GLSL_SVO_VOXEL_AIR)
    {
        crash(formatString("Sanity check failed: the air voxel was given the unexpected node ref $+$+", std::hex, airNodeRef));
    }
}

BlockAtlas::SpriteID BlockAtlas::addSprite(BlockSprite &sprite)
{
    // Check that the frame count is supported.
    size_t numFrames = sprite.frames.size();
    if (numFrames != 1 && numFrames != 2 && numFrames != 4 && numFrames != 8 && numFrames != 16)
    {
        crash(formatString(
            "BlockAtlas::addSprite() received a sprite with $+ frames; only 1, 2, 4, 8 or 16 frames are allowed.",
            numFrames
        ));

        sprite.frames.resize(16);
    }

    return makeSpriteNode(sprite, 3, glm::ivec3(0));
}

const std::vector<VoxelTree::Node>& BlockAtlas::getAtlasNodes()
{
    return atlasNodes;
}

size_t BlockAtlas::getNumVoxels()
{
    return finalVoxelCount;
}

void BlockAtlas::loadVoxelPalette(std::atomic<unsigned int> &tasksDone)
{
    voxelPalette[0].refract = 0x001A00FF;

    std::shared_ptr<Vk::Device> device = GraphicsContext::instance->getDevice();

    voxelPaletteDescPool = std::make_shared<Vk::DescriptorPool>(
        device,
        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        2,
        1
    );

    voxelPaletteDescSet = std::make_unique<Vk::DescriptorSet>(
        device,
        voxelPaletteDescPool,
        GraphicsContext::voxelPaletteDescriptorSetLayout
    );

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

    paletteStagingBuffer = std::make_shared<Vk::Buffer>(
        device,
        sizeof(Voxel) * voxelPalette.size(),
        VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
    );

    nodeSamplesStagingBuffer = std::make_shared<Vk::Buffer>(
        device,
        sizeof(NodeSample) * nodeSamples.size(),
        VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
    );

    memcpy(paletteStagingBuffer->hostdata(), voxelPalette.data(), sizeof(Voxel) * voxelPalette.size());
    memcpy(nodeSamplesStagingBuffer->hostdata(), nodeSamples.data(), sizeof(NodeSample) * nodeSamples.size());

    voxelPaletteBuffer = std::make_unique<Vk::Buffer>(
        device,
        sizeof(Voxel) * voxelPalette.size(),
        VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
        VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT
    );

    nodeSamplesBuffer = std::make_unique<Vk::Buffer>(
        device,
        sizeof(NodeSample) * nodeSamples.size(),
        VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
        VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT
    );

    cmdbuf->transitionBufferToTransferDst(*voxelPaletteBuffer);
    cmdbuf->copyBuffer(*paletteStagingBuffer, 0, *voxelPaletteBuffer, 0, sizeof(Voxel) * voxelPalette.size());
    cmdbuf->transitionBufferToShaderRead(*voxelPaletteBuffer);

    cmdbuf->transitionBufferToTransferDst(*nodeSamplesBuffer);
    cmdbuf->copyBuffer(*nodeSamplesStagingBuffer, 0, *nodeSamplesBuffer, 0, sizeof(NodeSample) * nodeSamples.size());
    cmdbuf->transitionBufferToShaderRead(*nodeSamplesBuffer);

    cmdbuf->end();

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

        voxelPaletteDescSet->setStorageBuffer(0, *voxelPaletteBuffer, 0, VK_WHOLE_SIZE);
        voxelPaletteDescSet->setStorageBuffer(1, *nodeSamplesBuffer, 0, VK_WHOLE_SIZE);

        paletteStagingBuffer = nullptr;
        finalVoxelCount = voxelPalette.size();
        voxelPalette = {};

        tasksDone++;
    });
}

Vk::DescriptorSet& BlockAtlas::getVoxelPaletteDescSet()
{
    return *voxelPaletteDescSet;
}