#include "Dx12ChunkData.h"

#include <DX12CommandQueue.h>

#include "Constants.h"

DX12ChunkData::DX12ChunkData(Framework::Point center)
    : ReferenceCounter(),
      aabbs(0),
      chunkInfoBuffer(0),
      blasScratchBuffer(0),
      blasBuffer(0),
      blockIndexBuffer(0),
      textureIdBuffer(0),
      chunkCenter(center),
      blockIndices(),
      textureIds(),
      changed(0),
      bufferChanged(0)
{}

DX12ChunkData::~DX12ChunkData()
{
    if (aabbs)
    {
        aabbs->release();
    }
    if (chunkInfoBuffer)
    {
        chunkInfoBuffer->release();
    }
    if (blasScratchBuffer)
    {
        blasScratchBuffer->release();
    }
    if (blasBuffer)
    {
        blasBuffer->release();
    }
    if (blockIndexBuffer)
    {
        blockIndexBuffer->release();
    }
    if (textureIdBuffer)
    {
        textureIdBuffer->release();
    }
}

Framework::DX12Buffer* DX12ChunkData::zChunkInfoBuffer() const
{
    return chunkInfoBuffer;
}

Framework::DX12Buffer* DX12ChunkData::zBlasBuffer() const
{
    return blasBuffer;
}

Framework::DX12Buffer* DX12ChunkData::zBlockIndexBuffer() const
{
    return blockIndexBuffer;
}

Framework::DX12Buffer* DX12ChunkData::zTextureIdBuffer() const
{
    return textureIdBuffer;
}

void DX12ChunkData::setBlock(int index, Block* zBlock)
{
    if (zBlock)
    {
        ArrayIterator<int> it = blockIndices.begin();
        ArrayIterator<int> textureIt = textureIds.begin();
        int pos = 0;
        while (it)
        {
            if (it.val() == index)
            {
                for (int i = 0; i < 6; ++i)
                {
                    textureIt.set(
                        zBlock->zTexture()->zPolygonTexture(i)->getId());
                    ++textureIt;
                }
                changed = 1;
                return;
            }
            ++it;
            ++pos;
            for (int i = 0; i < 6; ++i)
            {
                ++textureIt;
            }
        }
        blockIndices.add(index);
        for (int i = 0; i < 6; ++i)
        {
            textureIds.add(zBlock->zTexture()->zPolygonTexture(i)->getId());
        }
        changed = 1;
    }
    else
    {
        ArrayIterator<int> it = blockIndices.begin();
        ArrayIterator<int> textureIt = textureIds.begin();
        while (it)
        {
            if (it.val() == index)
            {
                for (int i = 0; i < 6; ++i)
                {
                    textureIt.remove();
                }
                it.remove();
                changed = 1;
                return;
            }
            ++it;
            for (int i = 0; i < 6; ++i)
            {
                ++textureIt;
            }
        }
    }
}

void DX12ChunkData::updateBuffers(
    ID3D12Device5* zDevice, Framework::DX12CommandQueue* zQueue)
{
    if (!aabbs)
    {
        aabbs = new DX12Buffer(sizeof(D3D12_RAYTRACING_AABB),
            zDevice,
            dynamic_cast<Framework::DX12CommandQueue*>(zQueue->getThis()),
            D3D12_RESOURCE_FLAG_NONE);
        aabbs->setLength(sizeof(D3D12_RAYTRACING_AABB));
        D3D12_RAYTRACING_AABB data = {chunkCenter.x - 8.f,
            chunkCenter.y - 8.f,
            0,
            chunkCenter.x + 8.f,
            chunkCenter.y + 8.f,
            WORLD_HEIGHT};
        aabbs->setData(&data, 1);
        aabbs->copyToGPU();
    }
    if (!blasBuffer)
    {
        D3D12_RAYTRACING_GEOMETRY_DESC geometryDesc = {};
        geometryDesc.Type
            = D3D12_RAYTRACING_GEOMETRY_TYPE_PROCEDURAL_PRIMITIVE_AABBS;
        geometryDesc.Flags = D3D12_RAYTRACING_GEOMETRY_FLAG_NONE;
        geometryDesc.AABBs.AABBCount = 1;
        geometryDesc.AABBs.AABBs.StartAddress
            = aabbs->zBuffer()->GetGPUVirtualAddress();
        geometryDesc.AABBs.AABBs.StrideInBytes = sizeof(D3D12_RAYTRACING_AABB);

        D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_INPUTS prebuildDesc;
        prebuildDesc.Type
            = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL;
        prebuildDesc.DescsLayout = D3D12_ELEMENTS_LAYOUT_ARRAY;
        prebuildDesc.NumDescs = 1;
        prebuildDesc.pGeometryDescs = &geometryDesc;
        prebuildDesc.Flags
            = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_BUILD_FLAG_NONE;

        D3D12_RAYTRACING_ACCELERATION_STRUCTURE_PREBUILD_INFO info = {};
        zDevice->GetRaytracingAccelerationStructurePrebuildInfo(
            &prebuildDesc, &info);

        blasBuffer = new DX12Buffer(1,
            zDevice,
            dynamic_cast<DX12CommandQueue*>(zQueue->getThis()),
            D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS);
        blasBuffer->setLength(
            ROUND_UP_POWER_OF_2(info.ResultDataMaxSizeInBytes, 256));
        blasBuffer->createBufferWithoutData(
            D3D12_RESOURCE_STATE_RAYTRACING_ACCELERATION_STRUCTURE);

        blasScratchBuffer = new DX12Buffer(1,
            zDevice,
            dynamic_cast<DX12CommandQueue*>(zQueue->getThis()),
            D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS);
        blasScratchBuffer->setLength(
            ROUND_UP_POWER_OF_2(info.ScratchDataSizeInBytes, 256));
        blasScratchBuffer->createBufferWithoutData(D3D12_RESOURCE_STATE_COMMON);

        D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_DESC buildDesc;
        buildDesc.Inputs.Type
            = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL;
        buildDesc.Inputs.DescsLayout = D3D12_ELEMENTS_LAYOUT_ARRAY;
        buildDesc.Inputs.NumDescs = 1;
        buildDesc.Inputs.pGeometryDescs = &geometryDesc;
        buildDesc.DestAccelerationStructureData
            = {blasBuffer->zBuffer()->GetGPUVirtualAddress()};
        buildDesc.ScratchAccelerationStructureData
            = {blasScratchBuffer->zBuffer()->GetGPUVirtualAddress()};
        buildDesc.SourceAccelerationStructureData = 0;
        buildDesc.Inputs.Flags
            = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_BUILD_FLAG_NONE;

        zQueue->zCommandList()->BuildRaytracingAccelerationStructure(
            &buildDesc, 0, nullptr);
    }
    if (blockIndexBuffer
        && blockIndexBuffer->getElementCount() < blockIndices.getEntryCount())
    {
        blockIndexBuffer->release();
        blockIndexBuffer = 0;
    }
    if (!blockIndexBuffer)
    {
        blockIndexBuffer = new DX12Buffer(sizeof(int),
            zDevice,
            dynamic_cast<DX12CommandQueue*>(zQueue->getThis()),
            D3D12_RESOURCE_FLAG_NONE);
        blockIndexBuffer->setLength(blockIndices.getEntryCount() * sizeof(int));
        bufferChanged = 1;
    }
    if (textureIdBuffer
        && textureIdBuffer->getElementCount() < textureIds.getEntryCount())
    {
        textureIdBuffer->release();
        textureIdBuffer = 0;
    }
    if (!textureIdBuffer)
    {
        textureIdBuffer = new DX12Buffer(sizeof(int),
            zDevice,
            dynamic_cast<DX12CommandQueue*>(zQueue->getThis()),
            D3D12_RESOURCE_FLAG_NONE);
        textureIdBuffer->setLength(textureIds.getEntryCount() * sizeof(int));
        bufferChanged = 1;
    }
    if (!chunkInfoBuffer)
    {
        chunkInfoBuffer = new DX12Buffer(sizeof(ChunkShaderInfo),
            zDevice,
            dynamic_cast<DX12CommandQueue*>(zQueue->getThis()),
            D3D12_RESOURCE_FLAG_NONE);
        chunkInfoBuffer->setLength(sizeof(ChunkShaderInfo));
    }
    if (changed)
    {
        int* blockIndexData = new int[blockIndices.getEntryCount()];
        int* textureIdData = new int[textureIds.getEntryCount()];
        ArrayIterator<int> it = blockIndices.begin();
        ArrayIterator<int> textureIt = textureIds.begin();
        int index = 0;
        while (it)
        {
            blockIndexData[index] = it.val();
            for (int i = 0; i < 6; i++)
            {
                textureIdData[index * 6 + i] = textureIt.val();
                ++textureIt;
            }
            ++it;
        }
        blockIndexBuffer->setData(blockIndexData, 1);
        textureIdBuffer->setData(textureIdData, 1);
        blockIndexBuffer->copyToGPU(sizeof(int) * blockIndices.getEntryCount());
        textureIdBuffer->copyToGPU(sizeof(int) * textureIds.getEntryCount());
        delete[] blockIndexData;
        delete[] textureIdData;
        ChunkShaderInfo info
            = {chunkCenter, (unsigned int)blockIndices.getEntryCount()};
        chunkInfoBuffer->setData(&info, 1);
        chunkInfoBuffer->copyToGPU();
    }
}