////////////////////////////////////////////////////////////
//~ DirectX12 libs

#pragma warning(push, 0)
# include <d3d12.h>
# include <dxgidebug.h>
# include <dxgi1_6.h>
#pragma warning(pop)

#pragma comment(lib, "d3d12")
#pragma comment(lib, "dxgi")

////////////////////////////////////////////////////////////
//~ Tweakable defines

#define GPU_D12_TearingIsAllowed        1
#define GPU_D12_FrameLatency            1
#define GPU_D12_SwapchainBufferCount    3
#define GPU_D12_SwapchainFlags          (((GPU_D12_TearingIsAllowed != 0)   * DXGI_SWAP_CHAIN_FLAG_ALLOW_TEARING) \
                                      | ((GPU_D12_FrameLatency != 0)        * DXGI_SWAP_CHAIN_FLAG_FRAME_LATENCY_WAITABLE_OBJECT))


#define GPU_D12_MaxCbvSrvUavDescriptors (1024 * 64)
#define GPU_D12_MaxSamplerDescriptors   (1024 * 1)
#define GPU_D12_MaxRtvDescriptors       (1024 * 1)

////////////////////////////////////////////////////////////
//~ Pipeline types

Struct(GPU_D12_PipelineDesc)
{
    VertexShader vs;
    PixelShader ps;
    ComputeShader cs;
    b32 is_wireframe;
    D3D12_PRIMITIVE_TOPOLOGY_TYPE topology_type;
    GPU_Format render_target_formats[GPU_MaxRasterTargets];
};

Struct(GPU_D12_Pipeline)
{
    GPU_D12_Pipeline *next_in_bin;
    u64 hash;

    GPU_D12_PipelineDesc desc;
    ID3D12PipelineState *pso;

    Fence ready_fence;
    b32 ok;
    String error;
};

Struct(GPU_D12_PipelineBin)
{
    Mutex mutex;
    GPU_D12_Pipeline *first;
};

////////////////////////////////////////////////////////////
//~ Descriptor types

Struct(GPU_D12_Descriptor)
{
    GPU_D12_Descriptor *next_free;
    struct GPU_D12_DescriptorHeap *heap;

    b32 valid;
    u32 index;
    D3D12_CPU_DESCRIPTOR_HANDLE handle;
} extern Readonly GPU_D12_NilDescriptor = {
    .index = U32Max
};

Struct(GPU_D12_DescriptorHeap)
{
    Arena *arena;

    D3D12_DESCRIPTOR_HEAP_TYPE type;
    u32 descriptor_size;
    ID3D12DescriptorHeap *d3d_heap;
    D3D12_CPU_DESCRIPTOR_HANDLE start_handle;

    Mutex mutex;
    GPU_D12_Descriptor *first_free;
    u32 allocated_count;
    u32 max_count;
};

////////////////////////////////////////////////////////////
//~ Resource types

Struct(GPU_D12_Resource)
{
    GPU_D12_Resource *next_free;
    GPU_ResourceDesc desc;

    ID3D12Resource *d3d_resource;
    D3D12_RESOURCE_STATES state;

    u64 buffer_size;  /* Actual size of buffer in GPU memory */

    GPU_D12_Descriptor *srv_descriptor;
    GPU_D12_Descriptor *uav_descriptor;
    GPU_D12_Descriptor *rtv_descriptor;
    GPU_D12_Descriptor *sampler_descriptor;

    u64 barrier_gen;
    D3D12_RESOURCE_BARRIER_TYPE barrier_type;
    D3D12_RESOURCE_STATES barrier_state_after;

    D3D12_GPU_VIRTUAL_ADDRESS buffer_gpu_address;
};

Struct(GPU_D12_ResourceReuseList)
{
    u64 hash;
    GPU_D12_ResourceReuseList *next;
    GPU_D12_ResourceReuseList *prev;
    GPU_D12_Resource *first;
};

Struct(GPU_D12_ResourceReuseListBin)
{
    Mutex mutex;
    GPU_D12_ResourceReuseList *first;
    GPU_D12_ResourceReuseList *last;
    GPU_D12_ResourceReuseList *first_free;
};

////////////////////////////////////////////////////////////
//~ Queue types

Struct(GPU_D12_Queue)
{
    GPU_D12_QueueDesc desc;
    ID3D12CommandQueue *d3d_queue;

    Mutex submit_mutex;
    ID3D12Fence *submit_fence;
    u64 submit_fence_target;
    struct GPU_D12_RawCommandList *first_submitted_cl;
    struct GPU_D12_RawCommandList *last_submitted_cl;

    Fence sync_fence;
};

////////////////////////////////////////////////////////////
//~ Raw command list types

Struct(GPU_D12_RawCommandList)
{
    GPU_D12_Queue *queue;
    GPU_D12_RawCommandList *next;

    u64 submit_fence_target;

    ID3D12CommandAllocator *ca;
    ID3D12GraphicsCommandList *cl;
};

////////////////////////////////////////////////////////////
//~ Command list types

Enum(GPU_D12_CommandKind)
{
    GPU_D12_CommandKind_None,

    /* Barrier */
    GPU_D12_CommandKind_TransitionToSrv,
    GPU_D12_CommandKind_TransitionToUav,
    GPU_D12_CommandKind_TransitionToRtv,
    GPU_D12_CommandKind_TransitionToCopySrc,
    GPU_D12_CommandKind_TransitionToCopyDst,
    GPU_D12_CommandKind_FlushUav,

    /* Copy */
    GPU_D12_CommandKind_Copy,

    /* Clear */
    GPU_D12_CommandKind_ClearRtv,

    /* Rasterize */
    GPU_D12_CommandKind_Rasterize,

    /* Compute */
    GPU_D12_CommandKind_Compute,
};

Struct(GPU_D12_Command)
{
    GPU_D12_Command *next;
    GPU_D12_CommandKind kind;
    union
    {
        struct
        {
            GPU_D12_Resource *resource;
            i32 rt_slot;
        } barrier;
        struct
        {
            GPU_D12_Resource *dst;
            GPU_D12_Resource *src;
            String src_string;
        } copy;
        struct
        {
            GPU_D12_Resource *resource;
        } clear;
        struct
        {
            u32 sig_size;
            u8 sig[256];
            VertexShader vs;
            PixelShader ps;
            u32 rts_count;
            GPU_Viewport viewport;
            GPU_Scissor scissor;
            u32 instances_count;
            GPU_D12_Resource *index_buffer;
            GPU_RasterizeMode mode;
        } rasterize;
        struct
        {
            u32 sig_size;
            u8 sig[256];
            ComputeShader cs;
            u32 num_threads_x;
            u32 num_threads_y;
            u32 num_threads_z;
        } compute;
    };
};

Struct(GPU_D12_CommandList)
{
    GPU_D12_CommandList *next;
    GPU_D12_Command *first;
    GPU_D12_Command *last;
    u64 count;

    GPU_QueueKind queue_kind;
};

////////////////////////////////////////////////////////////
//~ Swapchain types

Struct(GPU_D12_SwapchainBuffer)
{
    struct GPU_D12_Swapchain *swapchain;
    ID3D12Resource *d3d_resource;
    GPU_D12_Descriptor *rtv_descriptor;
    D3D12_RESOURCE_STATES state;
};

Struct(GPU_D12_Swapchain)
{
    GPU_D12_Swapchain *next;
    GPU_Format format;

    IDXGISwapChain3 *swapchain;
    HWND window_hwnd;
    HANDLE waitable;
    Vec2I32 resolution;

    GPU_D12_SwapchainBuffer buffers[GPU_D12_SwapchainBufferCount];
};

////////////////////////////////////////////////////////////
//~ State types

#define GPU_D12_NumResourceReuseBins 1024

Struct(GPU_D12_FiberState)
{
    GPU_D12_CommandList *first_free_command_list;
    GPU_D12_Command *first_free_command;
};

Struct(GPU_D12_SharedState)
{
    Atomic64Padded resource_barrier_gen;

    /* Stats */
    Atomic64 driver_resources_allocated;
    Atomic64 driver_descriptors_allocated;

    /* Queues */
    GPU_D12_Queue *queues[GPU_NumQueues];

    /* Rootsig */
    ID3D12RootSignature *bindless_rootsig;

    /* Pipelines */
    GPU_D12_PipelineBin pipeline_bins[1024];

    /* Descriptor heaps */
    GPU_D12_DescriptorHeap *cbv_srv_uav_heap;
    GPU_D12_DescriptorHeap *sampler_heap;
    GPU_D12_DescriptorHeap *rtv_heap;

    /* Resources */
    Mutex free_resources_mutex;
    GPU_D12_Resource *first_free_resource;
    GPU_D12_ResourceReuseListBin resource_reuse_bins[GPU_D12_NumResourceReuseBins];

    /* Swapchains */
    Mutex free_swapchains_mutex;
    GPU_D12_Swapchain *first_free_swapchain;

    /* Device */
    IDXGIFactory6 *factory;
    IDXGIAdapter3 *adapter;
    ID3D12Device *device;
} extern GPU_D12_shared_state;

////////////////////////////////////////////////////////////
//~ Helpers

GPU_D12_FiberState *GPU_D12_FiberStateFromId(i16 fiber_id);
DXGI_FORMAT GPU_D12_DxgiFormatFromGpuFormat(GPU_Format format);
GPU_D12_Command *GPU_D12_PushCmd(GPU_D12_CommandList *cl);
u64 GPU_D12_ReuseHashFromResourceDesc(GPU_ResourceDesc desc, u64 buffer_size);

////////////////////////////////////////////////////////////
//~ Startup

void GPU_D12_Startup(void);

////////////////////////////////////////////////////////////
//~ Initialization

//- Device initialization
void GPU_D12_InitDevice(void);

//- Queue initialization
JobDecl(GPU_D12_InitQueue, { GPU_D12_QueueDesc *descs; });

//- Heap initialization
GPU_D12_DescriptorHeap *GPU_D12_InitDescriptorHeap(D3D12_DESCRIPTOR_HEAP_TYPE type, D3D12_DESCRIPTOR_HEAP_FLAGS flags, u32 max_descs, u32 desc_size);

//- Rootsig initialization
void GPU_D12_InitRootsig(void);

////////////////////////////////////////////////////////////
//~ Pipeline operations

JobDecl(GPU_D12_LoadPipeline, { GPU_D12_Pipeline *pipeline; });
GPU_D12_Pipeline *GPU_D12_PipelineFromDesc(GPU_D12_PipelineDesc desc);

////////////////////////////////////////////////////////////
//~ Queue operations

GPU_D12_Queue *GPU_D12_QueueFromKind(GPU_QueueKind kind);

////////////////////////////////////////////////////////////
//~ Descriptor operations

GPU_D12_Descriptor *GPU_D12_AcquireDescriptor(GPU_D12_DescriptorHeap *heap);
void GPU_D12_ReleaseDescriptor(GPU_D12_Descriptor *descriptor);

////////////////////////////////////////////////////////////
//~ Raw command list operations

GPU_D12_RawCommandList *GPU_D12_BeginRawCommandList(GPU_QueueKind queue_kind);
u64 GPU_D12_EndRawCommandList(GPU_D12_RawCommandList *cl);

////////////////////////////////////////////////////////////
//~ Swapchain helpers

void GPU_D12_InitSwapchainResources(GPU_D12_Swapchain *swapchain);
GPU_D12_SwapchainBuffer *GPU_D12_UpdateSwapchain(GPU_D12_Swapchain *swapchain, Vec2I32 resolution);
i64 GPU_D12_BlitToSwapchain(GPU_D12_SwapchainBuffer *dst, GPU_D12_Resource *texture, Vec2I32 dst_p0, Vec2I32 dst_p1, Vec2I32 src_p0, Vec2I32 src_p1, Vec4 clear_color);

////////////////////////////////////////////////////////////
//~ Sync job

JobDecl(GPU_D12_StartQueueSync, EmptySig);