#include "CustomDX12API.h"

#include <d3d12.h>
#include <DX12Buffer.h>
#include <DX12Shader.h>

#include "CustomClosestHitShader.h"
#include "CustomMissShader.h"
#include "CustomRayGenShader.h"

using namespace Framework;

CustomDX12API::CustomDX12API()
    : DirectX12()
{}

void CustomDX12API::initializePipeline()
{
    DX12Shader* rayGenShader
        = new DX12Shader(CustomRayGenShader, sizeof(CustomRayGenShader));
    // RayGen from RayGen.hlsl
    DX12ShaderSignature* rayGenSignature = new DX12ShaderSignature();
    rayGenSignature->addRegisterUsageLinkedToDescriptorHeap(
        0, DX12_SHADER_REGISTER_U_UNORDERED_ACCESS, 0);
    rayGenSignature->addRegisterUsageLinkedToDescriptorHeap(
        1, DX12_SHADER_REGISTER_U_UNORDERED_ACCESS, 1);
    rayGenSignature->addRegisterUsageLinkedToDescriptorHeap(
        2, DX12_SHADER_REGISTER_B_CONST_BUFFER, 0);
    rayGenSignature->addRegisterUsageLinkedToDescriptorHeap(
        3, DX12_SHADER_REGISTER_T_SHADER_RESOURCE, 0);
    defaultRayGenerationShaderFunction = new DX12ShaderFunction(
        "RayGen", rayGenSignature, DX12_SHADER_FUNCTION_TYPE_RAY_GEN);
    rayGenShader->addFunction(defaultRayGenerationShaderFunction);
    pipeline->addShader(rayGenShader);

    DX12Shader* missShader
        = new DX12Shader(CustomMissShader, sizeof(CustomMissShader));
    // Miss from Miss.hlsl
    missShader->addFunction(new DX12ShaderFunction(
        "Miss", new DX12ShaderSignature(), DX12_SHADER_FUNCTION_TYPE_MISS));
    pipeline->addShader(missShader);

    DX12Shader* hitShader = new DX12Shader(
        CustomClosestHitShader, sizeof(CustomClosestHitShader));
    // ClosestHit from Hit.hlsl
    DX12ShaderFunction* closestHitFunction
        = new DX12ShaderFunction("ClosestHit",
            new DX12ShaderSignature(),
            DX12_SHADER_FUNCTION_TYPE_CLOSEST_HIT);
    hitShader->addFunction(closestHitFunction);
    pipeline->addShader(hitShader);

    defaultHitGroup = new DX12ShaderHitGroup("HitGroup");
    defaultHitGroup->setAttributeSize(8);
    defaultHitGroup->setPayloadSize(16);
    defaultHitGroup->setClosestHitShaderFunction(closestHitFunction);
    pipeline->addHitGroup(defaultHitGroup);

    pipeline->setMaxRecursionDepth(10);

    pipeline->createPipelineState(device, pfnD3D12SerializeRootSignature);
}

void CustomDX12API::initializeGlobalDescriptorHeap()
{
    DirectX12::initializeGlobalDescriptorHeap();
}

void CustomDX12API::fillShaderBindingTable(
    Framework::DX12ShaderBindingTable* zShaderBindingTable,
    Framework::Model3D* zModel,
    int objectIndex)
{
    DirectX12::fillShaderBindingTable(zShaderBindingTable, zModel, objectIndex);
}