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

#include <d3d12.h>
#include <dxgi1_4.h>
#include <directx/d3dx12.h>

#include <cassert>
#include <cstdio>

using namespace dx;

struct D3DSettings
{
    int width = 0;
    int height = 0;
};

class D3D
{
public:
    void Initialise(Window* window, const D3DSettings& settings)
    {
        m_window = window;
        m_settings = settings;

        UINT dxgiFactoryFlags = 0;
#ifdef DEBUG
        {
            ComPtr<ID3D12Debug> debug;
            D3D12GetDebugInterface(IID_PPV_ARGS(&debug));
            debug->EnableDebugLayer();
            dxgiFactoryFlags |= DXGI_CREATE_FACTORY_DEBUG;
        }
#endif
        ThrowIfFailed(CreateDXGIFactory2(
            dxgiFactoryFlags, IID_PPV_ARGS(&m_dxgi_factory)));

        // Prevent Alt+Enter from going into fullscreen.
        ThrowIfFailed(m_dxgi_factory->MakeWindowAssociation(
            m_window->GetWindowHandle(),
            DXGI_MWA_NO_ALT_ENTER));

        ThrowIfFailed(D3D12CreateDevice(
            /*pAdapter=*/nullptr, // Default adapter.
            D3D_FEATURE_LEVEL_11_0,
            IID_PPV_ARGS(&m_device)));

        m_rtv_descriptor_size = m_device->GetDescriptorHandleIncrementSize(
            D3D12_DESCRIPTOR_HEAP_TYPE_RTV);
        m_dsv_descriptor_size = m_device->GetDescriptorHandleIncrementSize(
            D3D12_DESCRIPTOR_HEAP_TYPE_DSV);
        m_cbv_descriptor_size = m_device->GetDescriptorHandleIncrementSize(
            D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV);

        const D3D12_COMMAND_QUEUE_DESC queue_desc =
        {
            .Type  = D3D12_COMMAND_LIST_TYPE_DIRECT,
            .Flags = D3D12_COMMAND_QUEUE_FLAG_NONE,
        };
        ThrowIfFailed(m_device->CreateCommandQueue(
            &queue_desc,
            IID_PPV_ARGS(&m_command_queue)));

        ThrowIfFailed(m_device->CreateCommandAllocator(
            queue_desc.Type,
            IID_PPV_ARGS(&m_command_allocator)));

        // The command allocator is the memory backing for the command list.
        // It is in the allocator's memory where the commands are stored.
        ThrowIfFailed(m_device->CreateCommandList(
            /*nodeMask=*/0,
            queue_desc.Type,
            m_command_allocator.Get(),
            /*pInitialState=*/nullptr,  // Pipeline state.
            IID_PPV_ARGS(&m_command_list)));

        // 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.
        ThrowIfFailed(m_command_list->Close());

        CreateDescriptorHeaps();

        CreateSwapChain();
        CreateSwapChainBufferRenderTargetViews();
        CreateDepthStencilBufferAndView();

        ThrowIfFailed(m_device->CreateFence(
            /*InitialValue=*/m_fence_value,
            D3D12_FENCE_FLAG_NONE,
            IID_PPV_ARGS(&m_fence)));

        if ((m_fence_event = CreateEvent(
            /*lpEventAttributes=*/nullptr,
            /*bManualReset=*/FALSE,
            /*bInitialState=*/FALSE,
            /*lpName=*/nullptr)) == 0)
        {
            ThrowIfFailed(HRESULT_FROM_WIN32(GetLastError()));
        }
    }

    void Render()
    {
        PopulateCommandList();

        ID3D12CommandList* command_lists[] = { m_command_list.Get() };
        m_command_queue->ExecuteCommandLists(
            _countof(command_lists), command_lists);

        ThrowIfFailed(m_swap_chain->Present(/*SyncInterval=*/1, /*Flags=*/0));
        m_current_back_buffer = m_swap_chain->GetCurrentBackBufferIndex();

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

private:
    void PopulateCommandList()
    {
        /// 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.
        ThrowIfFailed(m_command_allocator->Reset());

        // 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.
        ThrowIfFailed(m_command_list->Reset(
            m_command_allocator.Get(),
            /*pInitialState=*/nullptr));

        // Indicate that we intend to use the back buffer as a render target.
        const auto render_barrier = CD3DX12_RESOURCE_BARRIER::Transition(
            GetCurrentBackBuffer(),
            D3D12_RESOURCE_STATE_PRESENT,
            D3D12_RESOURCE_STATE_RENDER_TARGET);
        m_command_list->ResourceBarrier(1, &render_barrier);

        // Record commands.
        const float clear_colour[] = { 0.0f, 0.0f, 0.0f, 0.0f };
        m_command_list->ClearRenderTargetView(
            GetCurrentBackBufferView(),
            clear_colour,
            0,         // Number of rectangles in the following array.
            nullptr);  // No rectangles; clear the entire resource.

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

        // Indicate that we now intend to use the back buffer to present.
        const auto present_barrier = CD3DX12_RESOURCE_BARRIER::Transition(
            GetCurrentBackBuffer(),
            D3D12_RESOURCE_STATE_RENDER_TARGET,
            D3D12_RESOURCE_STATE_PRESENT);
        m_command_list->ResourceBarrier(1, &present_barrier);

        ThrowIfFailed(m_command_list->Close());
    }

    void WaitForPreviousFrame()
    {
        // Advance the fence value to mark commands up to this fence point.
        m_fence_value++;

        // The command queue will signal the new fence value when all commands
        // up to this point have finished execution.
        ThrowIfFailed(m_command_queue->Signal(m_fence.Get(), m_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 (m_fence->GetCompletedValue() < m_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 |m_fence_event| is to be fired when |m_fence|
            // reaches the |fence| value.
            ThrowIfFailed(m_fence->SetEventOnCompletion(
                m_fence_value, m_fence_event));

            WaitForSingleObject(m_fence_event, INFINITE);
        }
    }

    /// Creates RTV and DSV descriptor heaps.
    void CreateDescriptorHeaps()
    {
        assert(m_device);

        // The RTV heap must hold as many descriptors as we have buffers in the
        // swap chain.
        const D3D12_DESCRIPTOR_HEAP_DESC rtv_heap_desc =
        {
            .Type = D3D12_DESCRIPTOR_HEAP_TYPE_RTV,
            .NumDescriptors = SWAP_CHAIN_BUFFER_COUNT,
            .Flags = D3D12_DESCRIPTOR_HEAP_FLAG_NONE,
            .NodeMask = 0,
        };
        ThrowIfFailed(m_device->CreateDescriptorHeap(
            &rtv_heap_desc, IID_PPV_ARGS(&m_rtv_heap)));

        // For the depth/stencil buffer, we just need one view.
        const D3D12_DESCRIPTOR_HEAP_DESC dsv_heap_desc =
        {
            .Type = D3D12_DESCRIPTOR_HEAP_TYPE_DSV,
            .NumDescriptors = 1,
            .Flags = D3D12_DESCRIPTOR_HEAP_FLAG_NONE,
            .NodeMask = 0,
        };
        ThrowIfFailed(m_device->CreateDescriptorHeap(
            &dsv_heap_desc, IID_PPV_ARGS(&m_dsv_heap)));
    }

    /// Creates the application's swap chain.
    ///
    /// This method can be called multiple times to re-create the swap chain.
    void CreateSwapChain()
    {
        assert(m_dxgi_factory);
        assert(m_command_queue);

        SafeRelease(m_swap_chain);

        DXGI_SWAP_CHAIN_DESC1 desc =
        {
            .Width  = static_cast<UINT>(m_settings.width),
            .Height = static_cast<UINT>(m_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,
        };
        ComPtr<IDXGISwapChain1> swap_chain;
        ThrowIfFailed(m_dxgi_factory->CreateSwapChainForHwnd(
            m_command_queue.Get(),  // Swap chain uses queue to perform flush.
            m_window->GetWindowHandle(),
            &desc,
            /*pFullScreenDesc=*/nullptr,  // Running in windowed mode.
            /*pRestrictToOutput=*/nullptr,
            &swap_chain));
        ThrowIfFailed(swap_chain.As(&m_swap_chain));

        m_current_back_buffer = m_swap_chain->GetCurrentBackBufferIndex();
    }

    /// Creates RTVs for all of the swap chain's buffers.
    void CreateSwapChainBufferRenderTargetViews()
    {
        assert(m_device);
        assert(m_swap_chain);
        assert(m_rtv_heap);

        // Create the new buffer views.
        CD3DX12_CPU_DESCRIPTOR_HANDLE rtv_heap_handle(
            m_rtv_heap->GetCPUDescriptorHandleForHeapStart());
        for (int i = 0; i < SWAP_CHAIN_BUFFER_COUNT; ++i)
        {
            ThrowIfFailed(m_swap_chain->GetBuffer(
                i, IID_PPV_ARGS(&m_swap_chain_buffer[i])));

            m_device->CreateRenderTargetView(
                m_swap_chain_buffer[i].Get(), /*pDesc=*/nullptr, rtv_heap_handle);

            rtv_heap_handle.Offset(1, m_rtv_descriptor_size);
        }
    }

    /// Creates a depth/stencil buffer and its view.
    void CreateDepthStencilBufferAndView()
    {
        assert(m_device);

        const D3D12_RESOURCE_DESC depth_stencil_desc =
        {
            .Dimension = D3D12_RESOURCE_DIMENSION_TEXTURE2D,
            .Alignment = 0,
            .Width = static_cast<UINT64>(m_settings.width),
            .Height = static_cast<UINT>(m_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 =
        {
            .Format = depth_stencil_desc.Format,
            .DepthStencil = D3D12_DEPTH_STENCIL_VALUE
            {
                .Depth = 1.0f,
                .Stencil = 0,
            },
        };
        const CD3DX12_HEAP_PROPERTIES depth_stencil_heap_properties(
            D3D12_HEAP_TYPE_DEFAULT);
        ThrowIfFailed(m_device->CreateCommittedResource(
            &depth_stencil_heap_properties,
            D3D12_HEAP_FLAG_NONE,
            &depth_stencil_desc,
            D3D12_RESOURCE_STATE_COMMON,
            &opt_clear_value,
            IID_PPV_ARGS(&m_depth_stencil_buffer)));

        m_device->CreateDepthStencilView(
            m_depth_stencil_buffer.Get(),
            /*pDesc=*/nullptr,
            GetDepthStencilView());
    }

    ID3D12Resource* GetCurrentBackBuffer() const
    {
        return m_swap_chain_buffer[m_current_back_buffer].Get();
    }

    D3D12_CPU_DESCRIPTOR_HANDLE GetCurrentBackBufferView() const
    {
        assert(m_rtv_heap);
        assert(m_rtv_descriptor_size > 0);
        return CD3DX12_CPU_DESCRIPTOR_HANDLE(
            m_rtv_heap->GetCPUDescriptorHandleForHeapStart(),
            m_current_back_buffer,
            m_rtv_descriptor_size);
    }

    D3D12_CPU_DESCRIPTOR_HANDLE GetDepthStencilView() const
    {
        assert(m_dsv_heap);
        return m_dsv_heap->GetCPUDescriptorHandleForHeapStart();
    }

private:
    static constexpr int SWAP_CHAIN_BUFFER_COUNT = 2;  // Double-buffering.

    Window* m_window = nullptr;
    D3DSettings m_settings;

    ComPtr<IDXGIFactory4> m_dxgi_factory;
    ComPtr<ID3D12Device> m_device;

    ComPtr<ID3D12CommandQueue> m_command_queue;
    ComPtr<ID3D12CommandAllocator> m_command_allocator;
    ComPtr<ID3D12GraphicsCommandList> m_command_list;

    ComPtr<IDXGISwapChain3> m_swap_chain;

    ComPtr<ID3D12DescriptorHeap> m_rtv_heap;
    ComPtr<ID3D12DescriptorHeap> m_dsv_heap;

    ComPtr<ID3D12Resource> m_swap_chain_buffer[SWAP_CHAIN_BUFFER_COUNT];
    ComPtr<ID3D12Resource> m_depth_stencil_buffer;

    ComPtr<ID3D12Fence> m_fence;
    HANDLE m_fence_event = 0;
    UINT64 m_fence_value = 0;

    // Index to the buffer in the RTV descriptor heap that represents the
    // current back buffer.
    int m_current_back_buffer = 0;

    UINT m_rtv_descriptor_size = 0;
    UINT m_dsv_descriptor_size = 0;
    UINT m_cbv_descriptor_size = 0;
};

int main()
{
    try
    {
        const D3DSettings settings =
        {
            // TODO: use 960x600 or 1920x1200 depending on native resolution.
            .width = 1920,
            .height = 1200,
        };

        if (!WindowInitialise())
        {
            THROW("Failed to initialise the window subsystem");
        }
        {
            Window window;
            if (!window.Initialise(
                settings.width, settings.height, /*title=*/"D3D Application"))
            {
                THROW(GetWindowError());
            }

            D3D d3d;
            d3d.Initialise(&window, settings);

            while (!window.ShouldClose())
            {
                window.Update();
                d3d.Render();
                Sleep(10);
            }
        }
        WindowTerminate();

        return 0;
    }
    catch (const std::exception& e)
    {
        fprintf(stderr, "Exception caught: %s\n", e.what());
        return 1;
    }
}