#include "DX12Shader.h"

#include "Logging.h"

using namespace Framework;

Framework::DX12ShaderSignature::DX12ShaderSignature(ID3D12Device5* zDevice,
    PFN_D3D12_SERIALIZE_ROOT_SIGNATURE pfnD3D12SerializeRootSignature)
    : ReferenceCounter(),
      signature(0),
      zDevice(zDevice),
      pfnD3D12SerializeRootSignature(pfnD3D12SerializeRootSignature)
{}

Framework::DX12ShaderSignature::~DX12ShaderSignature()
{
    if (signature)
    {
        signature->Release();
    }
}

void Framework::DX12ShaderSignature::addRegisterUsage(
    DX12ShaderRegister registerType, int registerIndex)
{
    addRegisterUsage(registerType, registerIndex, 0);
}

void Framework::DX12ShaderSignature::addRegisterUsage(
    DX12ShaderRegister registerType, int registerIndex, int spaceIndex)
{
    registerUsages.add({registerType, registerIndex, spaceIndex});
}

void Framework::DX12ShaderSignature::createSignature()
{
    if (signature)
    {
        signature->Release();
        signature = 0;
    }
    D3D12_ROOT_PARAMETER descriptorTable;
    descriptorTable.ParameterType = D3D12_ROOT_PARAMETER_TYPE_DESCRIPTOR_TABLE;
    descriptorTable.ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL;
    descriptorTable.DescriptorTable.NumDescriptorRanges
        = registerUsages.getEntryCount();
    D3D12_DESCRIPTOR_RANGE* descriptorRanges
        = new D3D12_DESCRIPTOR_RANGE[registerUsages.getEntryCount()];
    int index = 0;
    for (const auto& usage : registerUsages)
    {
        D3D12_DESCRIPTOR_RANGE* range = &descriptorRanges[index];
        switch (usage.registerType)
        {
        case DX12_SHADER_REGISTER_B_CONST_BUFFER:
            range->RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_CBV;
            break;
        case DX12_SHADER_REGISTER_T_SHADER_RESOURCE:
            range->RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_SRV;
            break;
        case DX12_SHADER_REGISTER_U_UNORDERED_ACCESS:
            range->RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_UAV;
            break;
        default:
            Logging::error() << "Unknown register type for root signature: "
                             << usage.registerType;
            range->RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_SRV;
        }
        range->NumDescriptors
            = 1; // TODO: optimize when multiple descriptors are used in a row
                 // e.g. u0, u1, u2 with the same range then NumDescriptors can
                 // be 3 for u0 and u1 and u2 would not be needed in this array
        range->BaseShaderRegister = usage.registerIndex;
        range->RegisterSpace = usage.spaceIndex;
        range->OffsetInDescriptorsFromTableStart = index;
        index++;
    }
    descriptorTable.DescriptorTable.pDescriptorRanges = descriptorRanges;
    D3D12_ROOT_SIGNATURE_DESC rootDesc = {};
    rootDesc.NumParameters = 1;
    rootDesc.pParameters = &descriptorTable;
    rootDesc.Flags = D3D12_ROOT_SIGNATURE_FLAG_LOCAL_ROOT_SIGNATURE;
    ID3DBlob* pSigBlob = 0;
    ID3DBlob* pErrorBlob = 0;
    HRESULT hr = pfnD3D12SerializeRootSignature(
        &rootDesc, D3D_ROOT_SIGNATURE_VERSION_1_0, &pSigBlob, &pErrorBlob);
    if (pSigBlob)
    {
        zDevice->CreateRootSignature(0,
            pSigBlob->GetBufferPointer(),
            pSigBlob->GetBufferSize(),
            __uuidof(ID3D12RootSignature),
            (void**)&signature);
        pSigBlob->Release();
    }
    if (pErrorBlob)
    {
        std::string errorMessage(
            static_cast<const char*>(pErrorBlob->GetBufferPointer()),
            pErrorBlob->GetBufferSize());
        Logging::error() << "Failed to serialize root signature: "
                         << errorMessage;
        pErrorBlob->Release();
    }
    delete[] descriptorRanges;
}

ID3D12RootSignature* Framework::DX12ShaderSignature::zSignature() const
{
    return signature;
}