#include "DX12BLAS.h"

#include "DX12Buffer.h"
#include "DX12CommandQueue.h"

Framework::DX12BLAS::DX12BLAS(
    ID3D12Device5* zDevice, DX12DirectCommandQueue* zDirectQueue)
    : ReferenceCounter(),
      scratchBuffer(0),
      resultBuffer(0),
      zDevice(zDevice),
      zDirectQueue(zDirectQueue),
      geometryDesc(new D3D12_RAYTRACING_GEOMETRY_DESC())
{}

Framework::DX12BLAS::~DX12BLAS()
{
    if (scratchBuffer)
    {
        scratchBuffer->release();
    }
    if (resultBuffer)
    {
        resultBuffer->release();
    }
    delete geometryDesc;
}

void Framework::DX12BLAS::build(
    DX12Buffer* vertexBuffer, DX12Buffer* indexBuffer)
{
    memset(geometryDesc, 0, sizeof(D3D12_RAYTRACING_GEOMETRY_DESC));
    geometryDesc->Type = D3D12_RAYTRACING_GEOMETRY_TYPE_TRIANGLES;
    geometryDesc->Flags = D3D12_RAYTRACING_GEOMETRY_FLAG_NONE;
    geometryDesc->Triangles.VertexBuffer.StartAddress
        = vertexBuffer->zBuffer()->GetGPUVirtualAddress();
    geometryDesc->Triangles.VertexBuffer.StrideInBytes
        = vertexBuffer->getElementLength();
    geometryDesc->Triangles.VertexCount
        = (unsigned)vertexBuffer->getElementCount();
    geometryDesc->Triangles.VertexFormat = DXGI_FORMAT_R32G32B32_FLOAT;
    geometryDesc->Triangles.IndexBuffer
        = indexBuffer->zBuffer()->GetGPUVirtualAddress();
    geometryDesc->Triangles.IndexCount
        = (unsigned)indexBuffer->getElementCount();
    geometryDesc->Triangles.IndexFormat = DXGI_FORMAT_R32_UINT;

    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);

    scratchBuffer->setLength(
        ROUND_UP_POWER_OF_2((int)info.ScratchDataSizeInBytes, 256));
    scratchBuffer->createBufferWithoutData(D3D12_RESOURCE_STATE_COMMON);
    resultBuffer->setLength(
        ROUND_UP_POWER_OF_2((int)info.ResultDataMaxSizeInBytes, 256));
    resultBuffer->createBufferWithoutData(
        D3D12_RESOURCE_STATE_RAYTRACING_ACCELERATION_STRUCTURE);

    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
        = {resultBuffer->zBuffer()->GetGPUVirtualAddress()};
    buildDesc.ScratchAccelerationStructureData
        = {scratchBuffer->zBuffer()->GetGPUVirtualAddress()};
    buildDesc.SourceAccelerationStructureData = 0;
    buildDesc.Inputs.Flags
        = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_BUILD_FLAG_NONE;

    // Build the AS
    zDirectQueue->zCommandList()->BuildRaytracingAccelerationStructure(
        &buildDesc, 0, nullptr);
}

Framework::DX12Buffer* Framework::DX12BLAS::zScratchBuffer() const
{
    return scratchBuffer;
}

Framework::DX12Buffer* Framework::DX12BLAS::zResultBuffer() const
{
    return resultBuffer;
}