#include "pch.h"
#include "HC_Dx12.h"

#include "HC_Common.h"
#include "precompile/hudless_compare_PShader.h"
#include "precompile/hudless_compare_VShader.h"

#include <Config.h>

inline static int GetFormatGroup(DXGI_FORMAT format)
{
    switch (format)
    {

    case DXGI_FORMAT_R32G32B32A32_TYPELESS:
    case DXGI_FORMAT_R32G32B32A32_FLOAT:
    case DXGI_FORMAT_R32G32B32A32_UINT:
    case DXGI_FORMAT_R32G32B32A32_SINT:
        return 1;

    case DXGI_FORMAT_R32G32B32_TYPELESS:
    case DXGI_FORMAT_R32G32B32_FLOAT:
    case DXGI_FORMAT_R32G32B32_UINT:
    case DXGI_FORMAT_R32G32B32_SINT:
        return 2;

    case DXGI_FORMAT_R16G16B16A16_TYPELESS:
    case DXGI_FORMAT_R16G16B16A16_FLOAT:
    case DXGI_FORMAT_R16G16B16A16_UNORM:
    case DXGI_FORMAT_R16G16B16A16_UINT:
    case DXGI_FORMAT_R16G16B16A16_SNORM:
    case DXGI_FORMAT_R16G16B16A16_SINT:
        return 3;

    case DXGI_FORMAT_R10G10B10A2_TYPELESS:
    case DXGI_FORMAT_R10G10B10A2_UNORM:
    case DXGI_FORMAT_R10G10B10A2_UINT:
        return 4;

    case DXGI_FORMAT_R11G11B10_FLOAT:
        return 5;

    case DXGI_FORMAT_R8G8B8A8_TYPELESS:
    case DXGI_FORMAT_R8G8B8A8_UNORM:
    case DXGI_FORMAT_R8G8B8A8_UNORM_SRGB:
    case DXGI_FORMAT_R8G8B8A8_UINT:
    case DXGI_FORMAT_R8G8B8A8_SNORM:
    case DXGI_FORMAT_R8G8B8A8_SINT:
        return 6;

    case DXGI_FORMAT_B5G6R5_UNORM:
        return 7;

    case DXGI_FORMAT_B5G5R5A1_UNORM:
        return 8;

    case DXGI_FORMAT_B8G8R8A8_UNORM:
    case DXGI_FORMAT_B8G8R8A8_TYPELESS:
    case DXGI_FORMAT_B8G8R8A8_UNORM_SRGB:
        return 9;

    case DXGI_FORMAT_R10G10B10_XR_BIAS_A2_UNORM:
        return 10;

    case DXGI_FORMAT_B8G8R8X8_UNORM:
    case DXGI_FORMAT_B8G8R8X8_TYPELESS:
    case DXGI_FORMAT_B8G8R8X8_UNORM_SRGB:
        return 11;

    default:
        return -1;
    }
}

inline static bool CompareResourceFormats(DXGI_FORMAT sc, DXGI_FORMAT hudless)
{
    if (sc == hudless)
        return true;

    auto scGroup = GetFormatGroup(sc);
    auto hudlessGroup = GetFormatGroup(hudless);
    return scGroup == hudlessGroup;
}

bool HC_Dx12::CreateBufferResource(UINT index, ID3D12Device* InDevice, ID3D12Resource* InSource,
                                   D3D12_RESOURCE_STATES InState)
{
    LOG_DEBUG("[{0}] Start!", _name);

    auto resourceFlags = D3D12_RESOURCE_FLAG_ALLOW_RENDER_TARGET;

    auto result = Shader_Dx12::CreateBufferResource(InDevice, InSource, InState, &_buffer[index], resourceFlags);

    if (result)
    {
        _buffer[index]->SetName(L"HC_Buffer");
        _bufferState[index] = InState;
    }

    return result;
}

void HC_Dx12::ResourceBarrier(ID3D12GraphicsCommandList* cmdList, ID3D12Resource* resource,
                              D3D12_RESOURCE_STATES beforeState, D3D12_RESOURCE_STATES afterState)
{
    if (beforeState == afterState)
        return;

    D3D12_RESOURCE_BARRIER barrier = {};
    barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
    barrier.Transition.pResource = resource;
    barrier.Transition.StateBefore = beforeState;
    barrier.Transition.StateAfter = afterState;
    barrier.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
    cmdList->ResourceBarrier(1, &barrier);
}

void HC_Dx12::SetBufferState(UINT index, ID3D12GraphicsCommandList* InCommandList, D3D12_RESOURCE_STATES InState)
{
    if (_bufferState[index] == InState)
        return;

    ResourceBarrier(InCommandList, _buffer[index], _bufferState[index], InState);

    _bufferState[index] = InState;
}

HC_Dx12::HC_Dx12(std::string InName, ID3D12Device* InDevice) : Shader_Dx12(InName, InDevice)
{
    DXGI_SWAP_CHAIN_DESC scDesc {};
    if (State::Instance().currentSwapchain->GetDesc(&scDesc) != S_OK)
    {
        LOG_ERROR("Can't get swapchain desc!");
        return;
    }

    CD3DX12_STATIC_SAMPLER_DESC sampler(0);
    sampler.Filter = D3D12_FILTER_MIN_MAG_MIP_POINT;
    sampler.ShaderVisibility = D3D12_SHADER_VISIBILITY_PIXEL;
    sampler.AddressU = sampler.AddressV = sampler.AddressW = D3D12_TEXTURE_ADDRESS_MODE_CLAMP;

    if (!SetupRootSignature(InDevice, 2, 0, 1, 1, 0, 1, &sampler,
                            D3D12_ROOT_SIGNATURE_FLAG_ALLOW_INPUT_ASSEMBLER_INPUT_LAYOUT))
    {
        LOG_ERROR("Failed to setup root signature");
        return;
    }

    // Compile shaders
    UINT cflags = 0;
    ID3DBlob *vs, *ps;

    D3D12_GRAPHICS_PIPELINE_STATE_DESC graphicsPsoDesc {};
    graphicsPsoDesc.pRootSignature = _rootSignature;

    if (Config::Instance()->UsePrecompiledShaders.value_or_default())
    {
        graphicsPsoDesc.VS = CD3DX12_SHADER_BYTECODE(reinterpret_cast<const void*>(hudless_compare_VS_cso),
                                                     sizeof(hudless_compare_VS_cso));
        graphicsPsoDesc.PS = CD3DX12_SHADER_BYTECODE(reinterpret_cast<const void*>(hudless_compare_PS_cso),
                                                     sizeof(hudless_compare_PS_cso));
    }
    else
    {
        vs = CompileShader(hcCode.c_str(), "VSMain", "vs_5_1");
        if (vs != nullptr)
            graphicsPsoDesc.VS = { vs->GetBufferPointer(), vs->GetBufferSize() };
        else
            graphicsPsoDesc.VS = CD3DX12_SHADER_BYTECODE(reinterpret_cast<const void*>(hudless_compare_VS_cso),
                                                         sizeof(hudless_compare_VS_cso));

        ps = CompileShader(hcCode.c_str(), "PSMain", "ps_5_1");
        if (ps != nullptr)
            graphicsPsoDesc.PS = { ps->GetBufferPointer(), ps->GetBufferSize() };
        else
            graphicsPsoDesc.PS = CD3DX12_SHADER_BYTECODE(reinterpret_cast<const void*>(hudless_compare_PS_cso),
                                                         sizeof(hudless_compare_PS_cso));
    }

    graphicsPsoDesc.BlendState = CD3DX12_BLEND_DESC(D3D12_DEFAULT);
    graphicsPsoDesc.RasterizerState = CD3DX12_RASTERIZER_DESC(D3D12_DEFAULT);
    graphicsPsoDesc.DepthStencilState = CD3DX12_DEPTH_STENCIL_DESC(D3D12_DEFAULT);
    graphicsPsoDesc.DepthStencilState.DepthEnable = FALSE;
    graphicsPsoDesc.SampleMask = UINT_MAX;
    graphicsPsoDesc.PrimitiveTopologyType = D3D12_PRIMITIVE_TOPOLOGY_TYPE_TRIANGLE;
    graphicsPsoDesc.NumRenderTargets = 1;
    graphicsPsoDesc.RTVFormats[0] =
        Shader_Dx12::TranslateTypelessFormats(scDesc.BufferDesc.Format); // match swapchain RTV format (can be *_SRGB)
    graphicsPsoDesc.SampleDesc = { 1, 0 };

    auto result = InDevice->CreateGraphicsPipelineState(&graphicsPsoDesc, IID_PPV_ARGS(&_pipelineState));
    if (result != S_OK)
    {
        LOG_ERROR("CreateGraphicsPipelineState error: {:X}", (unsigned long) result);
        return;
    }

    // Create Constant Buffer
    D3D12_RESOURCE_DESC desc = CD3DX12_RESOURCE_DESC::Buffer(sizeof(InternalCompareParams));
    auto heapProps = CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_UPLOAD);

    result =
        InDevice->CreateCommittedResource(&heapProps, D3D12_HEAP_FLAG_NONE, &desc, D3D12_RESOURCE_STATE_GENERIC_READ,
                                          nullptr, IID_PPV_ARGS(&_constantBuffer));

    if (result != S_OK)
    {
        LOG_ERROR("CreateCommittedResource error {:X}", (unsigned int) result);
        return;
    }

    _init = InitHeaps(InDevice, _frameHeaps, HC_NUM_OF_HEAPS);
}

bool HC_Dx12::Dispatch(IDXGISwapChain3* sc, ID3D12GraphicsCommandList* cmdList, ID3D12Resource* hudless,
                       D3D12_RESOURCE_STATES state)
{
    if (sc == nullptr || hudless == nullptr || !_init)
        return false;

    DXGI_SWAP_CHAIN_DESC scDesc {};
    if (sc->GetDesc(&scDesc) != S_OK)
    {
        LOG_WARN("Can't get swapchain desc!");
        return false;
    }

    // Get SwapChain Buffer
    ID3D12Resource* scBuffer = nullptr;
    auto scIndex = sc->GetCurrentBackBufferIndex();
    auto result = sc->GetBuffer(scIndex, IID_PPV_ARGS(&scBuffer));

    if (result != S_OK)
    {
        LOG_ERROR("sc->GetBuffer({}) error: {:X}", scIndex, (unsigned long) result);
        return false;
    }

    scBuffer->Release();

    // Check Hudless Buffer
    D3D12_RESOURCE_DESC hudlessDesc = hudless->GetDesc();

    if (/*hudlessDesc.Format != scDesc.BufferDesc.Format ||*/ hudlessDesc.Width != scDesc.BufferDesc.Width ||
        hudlessDesc.Height != scDesc.BufferDesc.Height)
    {
        return false;
    }

    _counter++;
    _counter = _counter % HC_NUM_OF_HEAPS;

    if (!CreateBufferResource(_counter, _device, scBuffer, D3D12_RESOURCE_STATE_COPY_DEST))
    {
        LOG_ERROR("CreateBufferResource error!");
        return false;
    }

    // Copy Swapchain Buffer to read buffer
    SetBufferState(_counter, cmdList, D3D12_RESOURCE_STATE_COPY_DEST);
    ResourceBarrier(cmdList, scBuffer, D3D12_RESOURCE_STATE_PRESENT, D3D12_RESOURCE_STATE_COPY_SOURCE);

    if (_buffer[_counter] != nullptr)
        cmdList->CopyResource(_buffer[_counter], scBuffer);

    ResourceBarrier(cmdList, scBuffer, D3D12_RESOURCE_STATE_COPY_SOURCE, D3D12_RESOURCE_STATE_RENDER_TARGET);
    SetBufferState(_counter, cmdList, D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE);

    if (state != D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE)
        ResourceBarrier(cmdList, hudless, state, D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE);

    // Start setting pipeline
    UINT outWidth = scDesc.BufferDesc.Width;
    UINT outHeight = scDesc.BufferDesc.Height;

    FrameDescriptorHeap& currentHeap = _frameHeaps[_counter];

    // Create views
    // Fix for typeless formats, e.g. R8G8B8A8_TYPELESS, which can be used for both SRGB and UNORM
    if (CompareResourceFormats(scDesc.BufferDesc.Format, hudlessDesc.Format))
        CreateShaderResourceView(_device, hudless, currentHeap.GetSrvCPU(0), scDesc.BufferDesc.Format);
    else
        CreateShaderResourceView(_device, hudless, currentHeap.GetSrvCPU(0));

    CreateShaderResourceView(_device, _buffer[_counter], currentHeap.GetSrvCPU(1));
    CreateRenderTargetView(_device, scBuffer, currentHeap.GetRtvCPU(0), 0);

    InternalCompareParams constants {};
    constants.DiffThreshold = 0.003f;
    constants.PinkAmount = 0.6f;

    if (!CreateConstantsBuffer(_device, _constantBuffer, constants, currentHeap.GetCbvCPU(0)))
    {
        LOG_ERROR("[{0}] Failed to create a constants buffer", _name);
        return false;
    }

    ID3D12DescriptorHeap* heaps[] = { currentHeap.GetHeapCSU() };
    cmdList->SetDescriptorHeaps(_countof(heaps), heaps);

    cmdList->SetGraphicsRootSignature(_rootSignature);
    cmdList->SetPipelineState(_pipelineState);

    cmdList->SetGraphicsRootDescriptorTable(0, currentHeap.GetTableGPUStart());

    // Set RTV, viewport, scissor
    D3D12_CPU_DESCRIPTOR_HANDLE rtvHandles[] = { currentHeap.GetRtvCPU(0) };
    cmdList->OMSetRenderTargets(_countof(rtvHandles), rtvHandles, true, nullptr);

    D3D12_VIEWPORT vp {};
    vp.TopLeftX = 0.0f;
    vp.TopLeftY = 0.0f;
    vp.Width = static_cast<FLOAT>(outWidth);
    vp.Height = static_cast<FLOAT>(outHeight);
    vp.MinDepth = 0.0f;
    vp.MaxDepth = 1.0f;
    cmdList->RSSetViewports(1, &vp);

    D3D12_RECT rect { 0, 0, (LONG) outWidth, (LONG) outHeight };
    cmdList->RSSetScissorRects(1, &rect);

    // Fullscreen triangle
    cmdList->IASetPrimitiveTopology(D3D_PRIMITIVE_TOPOLOGY_TRIANGLELIST);
    cmdList->DrawInstanced(3, 1, 0, 0);

    ResourceBarrier(cmdList, scBuffer, D3D12_RESOURCE_STATE_RENDER_TARGET, D3D12_RESOURCE_STATE_PRESENT);

    if (state != D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE)
        ResourceBarrier(cmdList, hudless, D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE, state);

    return true;
}

HC_Dx12::~HC_Dx12()
{
    if (!_init || State::Instance().isShuttingDown)
        return;

    SAFE_RELEASE(_rootSignature);
    SAFE_RELEASE(_constantBuffer);

    for (int i = 0; i < HC_NUM_OF_HEAPS; i++)
    {
        _frameHeaps[i].ReleaseHeaps();
    }
}
