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

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

#define WIDTH  1920
#define HEIGHT 1200

#define SWAP_CHAIN_BUFFER_COUNT 2 // Double-buffering.
#define IMM_NUM_VERTS 1024

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

typedef struct D3D {
    Window*     pWindow;
    D3DSettings settings;

    IDXGIFactory4* pDxgiFactory;
    ID3D12Device*  pDevice;

    ID3D12CommandQueue* pCommandQueue;
    CommandRecorder     cmdRecPreamble;
    CommandRecorder     cmdRecPostamble;

    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;

    DxgImm* pImm;
} D3D;

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 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 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, DXGI_FORMAT swapChainRtvFormat, DXGI_SAMPLE_DESC swapChainSampleDesc) {
    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      = swapChainRtvFormat,
        .SampleDesc  = swapChainSampleDesc,
        .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);
    }
}

/// 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;
#ifndef NDEBUG
    ID3D12Debug* pDebug = NULL;
    TrapIfFailed(D3D12GetDebugInterface(&IID_ID3D12Debug, (&pDebug)));
    pDebug->lpVtbl->EnableDebugLayer(pDebug);
    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));

    TrapIfFailed(dxg_cmdrec_init(&d3d->cmdRecPreamble,  d3d->pDevice));
    TrapIfFailed(dxg_cmdrec_init(&d3d->cmdRecPostamble, d3d->pDevice));

    d3d_create_descriptor_heaps(d3d);

    const DXGI_FORMAT swapChainRtvFormat = DXGI_FORMAT_R8G8B8A8_UNORM;
    const DXGI_SAMPLE_DESC swapChainSampleDesc = {
        .Count   = 1,
        .Quality = 0,
    };

    d3d_create_swap_chain(d3d, swapChainRtvFormat, swapChainSampleDesc);
    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(NULL, FALSE, FALSE, NULL)) == NULL) {
        TrapIfFailed(HRESULT_FROM_WIN32(GetLastError()));
    }

    d3d->pImm = dxg_imm_init(
        d3d->pDevice, d3d->pCommandQueue, swapChainRtvFormat, swapChainSampleDesc, IMM_NUM_VERTS);
}

void d3d_destroy(D3D* d3d) {
    assert(d3d);
    dxg_imm_destroy(&d3d->pImm);
    dxg_cmdrec_destroy(&d3d->cmdRecPreamble);
    dxg_cmdrec_destroy(&d3d->cmdRecPostamble);
    SafeRelease(d3d->pCommandQueue);
    // Swap chain buffers are owned by the swap chain.
    SafeRelease(d3d->pDepthStencilBuffer);
    SafeRelease(d3d->pSwapChain);
    SafeRelease(d3d->pRtvHeap);
    SafeRelease(d3d->pDsvHeap);
    SafeRelease(d3d->pFence);
    SafeRelease(d3d->pDevice);
    SafeRelease(d3d->pDxgiFactory);
    CloseHandle(d3d->fence_event);
}

static void d3d_populate_command_lists(D3D* d3d) {
    assert(d3d);
    assert(d3d->cmdRecPreamble.pCmdAllocator);
    assert(d3d->cmdRecPreamble.pCmdList);
    assert(d3d->cmdRecPostamble.pCmdAllocator);
    assert(d3d->cmdRecPostamble.pCmdList);

    dxg_cmdrec_reset(&d3d->cmdRecPreamble);
    dxg_cmdrec_reset(&d3d->cmdRecPostamble);
    ID3D12GraphicsCommandList* pCmdListPre  = d3d->cmdRecPreamble.pCmdList;
    ID3D12GraphicsCommandList* pCmdListPost = d3d->cmdRecPostamble.pCmdList;

    // Indicate that we intend to use the back buffer as a render target and
    // depth/stencil for writing.
    const D3D12_RESOURCE_BARRIER render_barriers[] = {
        CD3DX12_RESOURCE_BARRIER_Transition(
            d3d_get_current_back_buffer(d3d),
            D3D12_RESOURCE_STATE_PRESENT,
            D3D12_RESOURCE_STATE_RENDER_TARGET),
        CD3DX12_RESOURCE_BARRIER_Transition(
            d3d->pDepthStencilBuffer,
            D3D12_RESOURCE_STATE_PRESENT,
            D3D12_RESOURCE_STATE_DEPTH_WRITE)
    };
    pCmdListPre->lpVtbl->ResourceBarrier(pCmdListPre, COUNTOF(render_barriers), render_barriers);

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

    // Clear the depth/stencil buffer.
    pCmdListPre->lpVtbl->ClearDepthStencilView(
        pCmdListPre,
        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 and depth/stencil
    // buffer to present.
    const D3D12_RESOURCE_BARRIER present_barriers[] = {
        CD3DX12_RESOURCE_BARRIER_Transition(
            d3d_get_current_back_buffer(d3d),
            D3D12_RESOURCE_STATE_RENDER_TARGET,
            D3D12_RESOURCE_STATE_PRESENT),
        CD3DX12_RESOURCE_BARRIER_Transition(
            d3d->pDepthStencilBuffer,
            D3D12_RESOURCE_STATE_DEPTH_WRITE,
            D3D12_RESOURCE_STATE_PRESENT)
    };
    pCmdListPost->lpVtbl->ResourceBarrier(pCmdListPost, COUNTOF(present_barriers), present_barriers);

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

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

    const D3D12_VIEWPORT viewport = {
        .TopLeftX = 0,
        .TopLeftY = 0,
        .Width    = WIDTH,
        .Height   = HEIGHT,
        .MinDepth = 0.0f,
        .MaxDepth = 1.0f
    };
    const D3D12_RECT scissor = {
        .bottom = HEIGHT,
        .left   = 0,
        .right  = WIDTH,
        .top    = 0,
    };
    const D3D12_CPU_DESCRIPTOR_HANDLE hBackBufferView   = d3d_get_current_back_buffer_view(d3d);
    const D3D12_CPU_DESCRIPTOR_HANDLE hDepthStencilView = d3d_get_depth_stencil_view(d3d);

    d3d_populate_command_lists(d3d); // Preamble and postamble.

    // Preamble.
    ID3D12CommandList* preamble[] = { (ID3D12CommandList*)d3d->cmdRecPreamble.pCmdList };
    d3d->pCommandQueue->lpVtbl->ExecuteCommandLists(
        d3d->pCommandQueue, COUNTOF(preamble), preamble);

    // Draw.
    const float o = 0.1f;
    const float triangle[] = {
        -1+o,  -1+o, 0,
           0,   1-o, 0,
        +1-o,  -1+o, 0,
    };
    dxg_imm_set_graphics_state(d3d->pImm, &viewport, hBackBufferView, hDepthStencilView);
    dxg_imm_draw_triangles(d3d->pImm, triangle, 1);
    dxg_imm_flush(d3d->pImm);

    // Postamble.
    ID3D12CommandList* postamble[] = { (ID3D12CommandList*)d3d->cmdRecPostamble.pCmdList };
    d3d->pCommandQueue->lpVtbl->ExecuteCommandLists(
        d3d->pCommandQueue, COUNTOF(postamble), postamble);

    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  = WIDTH,
        .height = HEIGHT,
    };

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

    d3d_destroy(&d3d);
    window_global_quit();
    return 0;
}