#include <dxg/dxcommon.h>
#include <dxwindow.h>

#include <assert.h>
#include <stdio.h>

#define SWAP_CHAIN_BUFFER_COUNT 2 // Double-buffering.

typedef struct D3DSettings
{
    int width;
    int height;
} D3DSettings;

typedef struct D3D {
    Window*     pWindow;
    D3DSettings settings;

    IDXGIFactory4* pDxgiFactory;
    ID3D12Device*  pDevice;

    ID3D12CommandQueue*        pCommandQueue;
    ID3D12CommandAllocator*    pCommandAllocator;
    ID3D12GraphicsCommandList* pCommandList;

    IDXGISwapChain3* pSwapChain;

    ID3D12DescriptorHeap* pRtvHeap;
    ID3D12DescriptorHeap* pDsvHeap;

    ID3D12Resource* pSwapChainBuffer[SWAP_CHAIN_BUFFER_COUNT];
    ID3D12Resource* pDepthStencilBuffer;

    ID3D12Fence* pFence;
    HANDLE       fence_event;
    UINT64       fence_value;

    UINT rtv_descriptor_size;
    UINT dsv_descriptor_size;
    UINT cbv_descriptor_size;
} D3D;

/// Creates the application's swap chain.
///
/// This method can be called multiple times to re-create the swap chain.
static void d3d_create_swap_chain(D3D* d3d) {
    assert(d3d);
    assert(d3d->pDxgiFactory);
    assert(d3d->pCommandQueue);
    assert(d3d->pWindow);

    SafeRelease(d3d->pSwapChain);

    DXGI_SWAP_CHAIN_DESC1 desc = (DXGI_SWAP_CHAIN_DESC1){
        .Width  = (UINT)(d3d->settings.width),
        .Height = (UINT)(d3d->settings.height),
        .Format = DXGI_FORMAT_R8G8B8A8_UNORM,
        .SampleDesc = (DXGI_SAMPLE_DESC){
            .Count   = 1,
            .Quality = 0,
        },
        .BufferUsage = DXGI_USAGE_RENDER_TARGET_OUTPUT,
        .BufferCount = SWAP_CHAIN_BUFFER_COUNT,
        .SwapEffect  = DXGI_SWAP_EFFECT_FLIP_DISCARD,
    };
    IDXGISwapChain1* pSwapChain = 0;
    TrapIfFailed(d3d->pDxgiFactory->lpVtbl->CreateSwapChainForHwnd(
        d3d->pDxgiFactory,
        (IUnknown*)d3d->pCommandQueue,  // Swap chain uses queue to perform flush.
        window_handle(d3d->pWindow),
        &desc,
        /*pFullScreenDesc=*/0,  // Running in windowed mode.
        /*pRestrictToOutput=*/0,
        &pSwapChain));
    //TrapIfFailed(pSwapChain.As(&d3d->pSwapChain));
    d3d->pSwapChain = (IDXGISwapChain3*)pSwapChain;
}

/// Creates RTVs for all of the swap chain's buffers.
static void d3d_create_swap_chain_buffer_render_target_views(D3D* d3d) {
    assert(d3d);
    assert(d3d->pDevice);
    assert(d3d->pSwapChain);
    assert(d3d->pSwapChainBuffer);
    assert(d3d->pRtvHeap);
    assert(d3d->rtv_descriptor_size > 0);

    // Create the new buffer views.
    D3D12_CPU_DESCRIPTOR_HANDLE rtv_heap_handle;
    d3d->pRtvHeap->lpVtbl->GetCPUDescriptorHandleForHeapStart(d3d->pRtvHeap, &rtv_heap_handle);
    for (int i = 0; i < SWAP_CHAIN_BUFFER_COUNT; ++i)
    {
        TrapIfFailed(d3d->pSwapChain->lpVtbl->GetBuffer(
            d3d->pSwapChain, i, &IID_ID3D12Resource, &d3d->pSwapChainBuffer[i]));

        d3d->pDevice->lpVtbl->CreateRenderTargetView(
            d3d->pDevice, d3d->pSwapChainBuffer[i], /*pDesc=*/0, rtv_heap_handle);

        rtv_heap_handle = OFFSET_HANDLE(rtv_heap_handle, 1, d3d->rtv_descriptor_size);
    }
}

static D3D12_CPU_DESCRIPTOR_HANDLE d3d_get_depth_stencil_view(D3D* d3d) {
    assert(d3d);
    assert(d3d->pDsvHeap);
    D3D12_CPU_DESCRIPTOR_HANDLE handle;
    d3d->pDsvHeap->lpVtbl->GetCPUDescriptorHandleForHeapStart(d3d->pDsvHeap, &handle);
    return handle;
}

/// Creates a depth/stencil buffer and its view.
static void d3d_create_depth_stencil_buffer_and_view(D3D* d3d) {
    assert(d3d);
    assert(d3d->pDevice);

    const D3D12_RESOURCE_DESC depth_stencil_desc = (D3D12_RESOURCE_DESC){
        .Dimension        = D3D12_RESOURCE_DIMENSION_TEXTURE2D,
        .Alignment        = 0,
        .Width            = (UINT64)(d3d->settings.width),
        .Height           = (UINT)(d3d->settings.height),
        .DepthOrArraySize = 1,
        .MipLevels        = 1,
        .Format           = DXGI_FORMAT_D24_UNORM_S8_UINT,
        .SampleDesc       = (DXGI_SAMPLE_DESC){
            .Count   = 1,
            .Quality = 0,
        },
        .Layout = D3D12_TEXTURE_LAYOUT_UNKNOWN,
        .Flags  = D3D12_RESOURCE_FLAG_ALLOW_DEPTH_STENCIL,
    };
    const D3D12_CLEAR_VALUE opt_clear_value = (D3D12_CLEAR_VALUE){
        .Format       = depth_stencil_desc.Format,
        .DepthStencil = (D3D12_DEPTH_STENCIL_VALUE){
            .Depth    = 1.0f,
            .Stencil  = 0,
        },
    };
    const D3D12_HEAP_PROPERTIES depth_stencil_heap_properties = (D3D12_HEAP_PROPERTIES){
        .Type                 = D3D12_HEAP_TYPE_DEFAULT,
        .CPUPageProperty      = D3D12_CPU_PAGE_PROPERTY_UNKNOWN,
        .MemoryPoolPreference = D3D12_MEMORY_POOL_UNKNOWN,
        .CreationNodeMask     = 1,
        .VisibleNodeMask      = 1,
    };
    TrapIfFailed(d3d->pDevice->lpVtbl->CreateCommittedResource(
        d3d->pDevice,
        &depth_stencil_heap_properties,
        D3D12_HEAP_FLAG_NONE,
        &depth_stencil_desc,
        D3D12_RESOURCE_STATE_COMMON,
        &opt_clear_value,
        &IID_ID3D12Resource,
        &d3d->pDepthStencilBuffer));

    d3d->pDevice->lpVtbl->CreateDepthStencilView(
        d3d->pDevice,
        d3d->pDepthStencilBuffer,
        /*pDesc=*/0,
        d3d_get_depth_stencil_view(d3d));
}

/// Creates RTV and DSV descriptor heaps.
static void d3d_create_descriptor_heaps(D3D* d3d) {
    assert(d3d);
    assert(d3d->pDevice);

    // The RTV heap must hold as many descriptors as we have buffers in the
    // swap chain.
    const D3D12_DESCRIPTOR_HEAP_DESC rtv_heap_desc  = (D3D12_DESCRIPTOR_HEAP_DESC){
        .Type           = D3D12_DESCRIPTOR_HEAP_TYPE_RTV,
        .NumDescriptors = SWAP_CHAIN_BUFFER_COUNT,
        .Flags          = D3D12_DESCRIPTOR_HEAP_FLAG_NONE,
        .NodeMask       = 0,
    };
    TrapIfFailed(d3d->pDevice->lpVtbl->CreateDescriptorHeap(
        d3d->pDevice, &rtv_heap_desc, &IID_ID3D12DescriptorHeap, &d3d->pRtvHeap));

    // For the depth/stencil buffer, we just need one view.
    const D3D12_DESCRIPTOR_HEAP_DESC dsv_heap_desc = (D3D12_DESCRIPTOR_HEAP_DESC){
        .Type           = D3D12_DESCRIPTOR_HEAP_TYPE_DSV,
        .NumDescriptors = 1,
        .Flags          = D3D12_DESCRIPTOR_HEAP_FLAG_NONE,
        .NodeMask       = 0,
    };
    TrapIfFailed(d3d->pDevice->lpVtbl->CreateDescriptorHeap(
        d3d->pDevice, &dsv_heap_desc, &IID_ID3D12DescriptorHeap, &d3d->pDsvHeap));
}

static void d3d_init(D3D* d3d, Window* pWindow, const D3DSettings* pSettings) {
    assert(d3d);
    assert(pWindow);
    assert(pSettings);

    d3d->pWindow  = pWindow;
    d3d->settings = *pSettings;

    UINT dxgiFactoryFlags = 0;
#ifdef DEBUG
    ID3D12Debug* debug = 0;
    D3D12GetDebugInterface(&IID_ID3D12Debug, (&debug));
    debug->EnableDebugLayer();
    dxgiFactoryFlags |= DXGI_CREATE_FACTORY_DEBUG;
#endif
    TrapIfFailed(CreateDXGIFactory2(
        dxgiFactoryFlags, &IID_IDXGIFactory4, &d3d->pDxgiFactory));

    // Prevent Alt+Enter from going into fullscreen.
    TrapIfFailed(d3d->pDxgiFactory->lpVtbl->MakeWindowAssociation(
        d3d->pDxgiFactory,
        window_handle(d3d->pWindow),
        DXGI_MWA_NO_ALT_ENTER));

    TrapIfFailed(D3D12CreateDevice(
        /*pAdapter=*/0, // Default adapter.
        D3D_FEATURE_LEVEL_11_0,
        &IID_ID3D12Device,
        &d3d->pDevice));

    d3d->rtv_descriptor_size = d3d->pDevice->lpVtbl->GetDescriptorHandleIncrementSize(
        d3d->pDevice,
        D3D12_DESCRIPTOR_HEAP_TYPE_RTV);
    d3d->dsv_descriptor_size = d3d->pDevice->lpVtbl->GetDescriptorHandleIncrementSize(
        d3d->pDevice,
        D3D12_DESCRIPTOR_HEAP_TYPE_DSV);
    d3d->cbv_descriptor_size = d3d->pDevice->lpVtbl->GetDescriptorHandleIncrementSize(
        d3d->pDevice,
        D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV);

    const D3D12_COMMAND_QUEUE_DESC queue_desc = (D3D12_COMMAND_QUEUE_DESC){
        .Type  = D3D12_COMMAND_LIST_TYPE_DIRECT,
        .Flags = D3D12_COMMAND_QUEUE_FLAG_NONE,
    };
    TrapIfFailed(d3d->pDevice->lpVtbl->CreateCommandQueue(
        d3d->pDevice,
        &queue_desc,
        &IID_ID3D12CommandQueue,
        &d3d->pCommandQueue));

    // The command allocator is the memory backing for the command list.
    // It is in the allocator's memory where the commands are stored.
    TrapIfFailed(d3d->pDevice->lpVtbl->CreateCommandAllocator(
        d3d->pDevice,
        queue_desc.Type,
        &IID_ID3D12CommandAllocator,
        &d3d->pCommandAllocator));

    TrapIfFailed(d3d->pDevice->lpVtbl->CreateCommandList(
        d3d->pDevice,
        /*nodeMask=*/0,
        queue_desc.Type,
        d3d->pCommandAllocator,
        /*pInitialState=*/0,  // Pipeline state.
        &IID_ID3D12CommandList,
        &d3d->pCommandList));

    // Command lists are in the "open" state after they are created. It is
    // easier to assume that they start in the "closed" state at each
    // iteration of the main loop, however. The Reset() method, which we'll
    // use later, also expects the command list to be closed.
    TrapIfFailed(d3d->pCommandList->lpVtbl->Close(d3d->pCommandList));

    d3d_create_descriptor_heaps(d3d);

    d3d_create_swap_chain(d3d);
    d3d_create_swap_chain_buffer_render_target_views(d3d);
    d3d_create_depth_stencil_buffer_and_view(d3d);

    TrapIfFailed(d3d->pDevice->lpVtbl->CreateFence(
        d3d->pDevice,
        d3d->fence_value,
        D3D12_FENCE_FLAG_NONE,
        &IID_ID3D12Fence,
        &d3d->pFence));

    if ((d3d->fence_event = CreateEvent(0, FALSE, FALSE, 0)) == 0) {
        TrapIfFailed(HRESULT_FROM_WIN32(GetLastError()));
    }
}

static D3D12_CPU_DESCRIPTOR_HANDLE d3d_get_current_back_buffer_view(const D3D* d3d) {
    assert(d3d);
    assert(d3d->pSwapChain);
    assert(d3d->pRtvHeap);
    assert(d3d->rtv_descriptor_size > 0);
    D3D12_CPU_DESCRIPTOR_HANDLE rtv_handle;
    d3d->pRtvHeap->lpVtbl->GetCPUDescriptorHandleForHeapStart(d3d->pRtvHeap, &rtv_handle);
    return CD3DX12_CPU_DESCRIPTOR_HANDLE(
        rtv_handle,
        d3d->pSwapChain->lpVtbl->GetCurrentBackBufferIndex(d3d->pSwapChain),
        d3d->rtv_descriptor_size);
}

static ID3D12Resource* d3d_get_current_back_buffer(const D3D* d3d) {
    assert(d3d);
    assert(d3d->pSwapChain);
    assert(d3d->pSwapChainBuffer);
    return d3d->pSwapChainBuffer[d3d->pSwapChain->lpVtbl->GetCurrentBackBufferIndex(d3d->pSwapChain)];
}

static void d3d_populate_command_list(D3D* d3d) {
    assert(d3d);
    assert(d3d->pCommandAllocator);
    assert(d3d->pCommandList);

    /// Note that we skip the following two items:
    ///
    ///   1. RSSetViewports()
    ///   2. OMSetRenderTargets()
    ///
    /// This application does not render anything useful, it simply clears
    /// the back buffer and depth/stencil view. Clearing both resources
    /// does not require a viewport to be set or the OM (output-merger
    /// stage) to be configured.

    // A command allocator can only be reset when its associated command
    // lists are finished executing on the GPU. This requires
    // synchronisation.
    TrapIfFailed(d3d->pCommandAllocator->lpVtbl->Reset(d3d->pCommandAllocator));

    // A command list can be reset as soon as it is executed with
    // ExecuteCommandList(). Reset() does require that the command list is
    // in a "closed" state, however, which is why we Close() it right away
    // after creation.
    TrapIfFailed(d3d->pCommandList->lpVtbl->Reset(
        d3d->pCommandList,
        d3d->pCommandAllocator,
        /*pInitialState=*/0));

    // Indicate that we intend to use the back buffer as a render target.
    const D3D12_RESOURCE_BARRIER render_barrier = CD3DX12_RESOURCE_BARRIER_Transition(
        d3d_get_current_back_buffer(d3d),
        D3D12_RESOURCE_STATE_PRESENT,
        D3D12_RESOURCE_STATE_RENDER_TARGET);
    d3d->pCommandList->lpVtbl->ResourceBarrier(d3d->pCommandList, 1, &render_barrier);

    // Record commands.
    const float clear_colour[] = { 0.0f, 0.502f, 0.494f, 0.0f };
    d3d->pCommandList->lpVtbl->ClearRenderTargetView(
        d3d->pCommandList,
        d3d_get_current_back_buffer_view(d3d),
        clear_colour,
        0,  // Number of rectangles in the following array.
        0); // No rectangles; clear the entire resource.

    d3d->pCommandList->lpVtbl->ClearDepthStencilView(
        d3d->pCommandList,
        d3d_get_depth_stencil_view(d3d),
        D3D12_CLEAR_FLAG_DEPTH | D3D12_CLEAR_FLAG_STENCIL,
        1.0f, // Depth.
        0,    // Stencil.
        0,    // Number of rectangles in the following array.
        0);   // No rectangles; clear the entire resource view.

    // Indicate that we now intend to use the back buffer to present.
    const D3D12_RESOURCE_BARRIER present_barrier = CD3DX12_RESOURCE_BARRIER_Transition(
        d3d_get_current_back_buffer(d3d),
        D3D12_RESOURCE_STATE_RENDER_TARGET,
        D3D12_RESOURCE_STATE_PRESENT);
    d3d->pCommandList->lpVtbl->ResourceBarrier(d3d->pCommandList, 1, &present_barrier);

    // A command list must be closed before it can be executed.
    TrapIfFailed(d3d->pCommandList->lpVtbl->Close(d3d->pCommandList));
}

static void d3d_wait_for_previous_frame(D3D* d3d) {
    assert(d3d);
    assert(d3d->pFence);
    assert(d3d->pCommandQueue);

    // Advance the fence value to mark commands up to this fence point.
    d3d->fence_value++;

    // The command queue will signal the new fence value when all commands
    // up to this point have finished execution.
    TrapIfFailed(d3d->pCommandQueue->lpVtbl->Signal(
        d3d->pCommandQueue, d3d->pFence, d3d->fence_value));

    // Wait for commands to finish execution.
    // It is possible that execution has already finished by the time we
    // get here, so first check the fence's completed value.
    if (d3d->pFence->lpVtbl->GetCompletedValue(d3d->pFence) < d3d->fence_value) {
        // Commands are still being executed. Configure a Windows event
        // and wait for it. The event fires when the commands have finished
        // execution.

        // Indicate that |fence_event| is to be fired when |fence|
        // reaches the new fence value.
        TrapIfFailed(d3d->pFence->lpVtbl->SetEventOnCompletion(
            d3d->pFence, d3d->fence_value, d3d->fence_event));

        // Will wake up when the fence takes on the new fence value.
        WaitForSingleObject(d3d->fence_event, INFINITE);
    }
}

static void d3d_render(D3D* d3d) {
    assert(d3d);
    assert(d3d->pCommandQueue);
    assert(d3d->pSwapChain);

    d3d_populate_command_list(d3d);

    ID3D12CommandList* command_lists[] = { (ID3D12CommandList*)d3d->pCommandList };
    d3d->pCommandQueue->lpVtbl->ExecuteCommandLists(
        d3d->pCommandQueue, _countof(command_lists), command_lists);

    TrapIfFailed(d3d->pSwapChain->lpVtbl->Present(
        d3d->pSwapChain, /*SyncInterval=*/1, /*Flags=*/0));

    // It is not efficient to wait for the frame to complete here, but it
    // is simple and sufficient for this application.
    d3d_wait_for_previous_frame(d3d);
}

int main(int argc, const char** argv) {
    const D3DSettings settings = (D3DSettings){
        .width  = 1920,
        .height = 1200,
    };

    if (!window_global_init()) {
        TRAP("Failed to initialise the window subsystem");
    }

    Window* window = window_init(settings.width, settings.height, "D3D Application");
    if (!window) {
        TRAP(window_get_error());
    }

    D3D d3d = {0};
    d3d_init(&d3d, window, &settings);

    while (!window_should_close(window))
    {
        window_update(window);
        d3d_render(&d3d);
        Sleep(10);
    }

    window_global_quit();
    return 0;
}