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

#include <d3d12.h>
#include <dxgidebug.h>
#include <dxgi1_6.h>

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

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

#define GPU_D12_TearingIsAllowed        1
#define GPU_D12_FrameLatency            0  /* TODO: Set this to 1 */
#define GPU_D12_SwapchainBufferCount    4
#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;
    D3D12_PRIMITIVE_TOPOLOGY_TYPE topology_type;
    GPU_Format render_target_formats[GPU_MaxRenderTargets];
};

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_CpuDescriptor)
{
    GPU_D12_CpuDescriptor *next_free;
    struct GPU_D12_CpuDescriptorHeap *heap;

    u32 index;
    D3D12_CPU_DESCRIPTOR_HANDLE handle;
};

Struct(GPU_D12_CpuDescriptorHeap)
{
    Arena *arena;

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

    Mutex mutex;
    GPU_D12_CpuDescriptor *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 reuse_hash;

    GPU_D12_CpuDescriptor *srv_descriptor;
    GPU_D12_CpuDescriptor *uav_descriptor;
    GPU_D12_CpuDescriptor *rtv_descriptor;
    GPU_D12_CpuDescriptor *sampler_descriptor;

    D3D12_GPU_VIRTUAL_ADDRESS buffer_gpu_address;
};

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

Struct(GPU_D12_QueueDesc)
{
    GPU_QueueKind kind;
    D3D12_COMMAND_LIST_TYPE d3d_type;
    D3D12_COMMAND_QUEUE_PRIORITY d3d_priority;
    String dbg_name;
};

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_Clear,

    /* 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;
            Vec4 value;
        } 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;

    b32 has_rasterize_cmd;
    b32 has_compute_cmd;
};

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

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

Struct(GPU_D12_Swapchain)
{
    GPU_D12_Swapchain *next;

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

    GPU_D12_SwapchainBuffer buffers[GPU_D12_SwapchainBufferCount];
};

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

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

Struct(GPU_D12_SharedState)
{
    GPU_D12_FiberState *fiber_states[MaxFibers];

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

    /* Rootsig */
    ID3D12RootSignature *bindless_rootsig;

    /* Pipelines */
    GPU_D12_PipelineBin pipeline_bins[1024];

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

    /* Resources */
    Mutex free_resources_mutex;
    GPU_D12_Resource *first_free_resource;

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

    /* Device */
    IDXGIFactory6 *factory;
    IDXGIAdapter1 *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);

////////////////////////////////
//~ 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_CpuDescriptorHeap *GPU_D12_InitCpuDescriptorHeap(D3D12_DESCRIPTOR_HEAP_TYPE type, 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_CpuDescriptor *GPU_D12_AcquireCpuDescriptor(GPU_D12_CpuDescriptorHeap *heap);
void GPU_D12_ReleaseCpuDescriptor(GPU_D12_CpuDescriptor *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 *swapchain_buffer, GPU_D12_Resource *texture);

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

JobDecl(GPU_D12_StartQueueSync, EmptySig);