#include "renderer/renderer.h"

#include <d3d12.h>
#include <dxgi1_5.h>

#include "imgui.h"
#include "imgui_impl_win32.h"
#include "imgui_impl_dx12.h"

#ifdef _DEBUG
#define DX12_ENABLE_DEBUG_LAYER
#endif

#ifdef DX12_ENABLE_DEBUG_LAYER
#include <dxgidebug.h>
#pragma comment(lib, "dxguid.lib")
#endif

#define NUM_BACK_BUFFERS  2
#define SRV_HEAP_SIZE     64

struct FrameContext {
    ID3D12CommandAllocator *command_allocator;
    UINT64                  fence_value;
};

struct SrvHeapAllocator {
    ID3D12DescriptorHeap       *heap;
    D3D12_DESCRIPTOR_HEAP_TYPE  heap_type;
    D3D12_CPU_DESCRIPTOR_HANDLE start_cpu;
    D3D12_GPU_DESCRIPTOR_HANDLE start_gpu;
    UINT                        increment;
    int                         free_indices[SRV_HEAP_SIZE];
    int                         free_count;
};

static void srv_alloc_init(SrvHeapAllocator *a, ID3D12Device *device, ID3D12DescriptorHeap *heap)
{
    a->heap      = heap;
    a->heap_type = heap->GetDesc().Type;
    a->start_cpu = heap->GetCPUDescriptorHandleForHeapStart();
    a->start_gpu = heap->GetGPUDescriptorHandleForHeapStart();
    a->increment = device->GetDescriptorHandleIncrementSize(a->heap_type);

    D3D12_DESCRIPTOR_HEAP_DESC desc = heap->GetDesc();
    a->free_count = (int)desc.NumDescriptors;
    for (int i = 0; i < a->free_count; i++)
        a->free_indices[i] = a->free_count - 1 - i;
}

static void srv_alloc_alloc(SrvHeapAllocator *a, D3D12_CPU_DESCRIPTOR_HANDLE *out_cpu, D3D12_GPU_DESCRIPTOR_HANDLE *out_gpu)
{
    int idx = a->free_indices[--a->free_count];
    out_cpu->ptr = a->start_cpu.ptr + (idx * a->increment);
    out_gpu->ptr = a->start_gpu.ptr + (idx * a->increment);
}

static void srv_alloc_free(SrvHeapAllocator *a, D3D12_CPU_DESCRIPTOR_HANDLE cpu, D3D12_GPU_DESCRIPTOR_HANDLE gpu)
{
    (void)gpu;
    int idx = (int)((cpu.ptr - a->start_cpu.ptr) / a->increment);
    a->free_indices[a->free_count++] = idx;
}

struct Renderer {
    HWND                         hwnd;
    int32_t                      width;
    int32_t                      height;
    int32_t                      frame_count;
    UINT                         frame_index;

    ID3D12Device                *device;
    ID3D12CommandQueue          *command_queue;
    IDXGISwapChain3             *swap_chain;
    HANDLE                       swap_chain_waitable;
    bool                         swap_chain_occluded;
    bool                         tearing_support;

    ID3D12DescriptorHeap        *rtv_heap;
    ID3D12DescriptorHeap        *srv_heap;
    SrvHeapAllocator             srv_alloc;

    FrameContext                 frames[NUM_BACK_BUFFERS];
    ID3D12Resource              *render_targets[NUM_BACK_BUFFERS];
    D3D12_CPU_DESCRIPTOR_HANDLE  rtv_descriptors[NUM_BACK_BUFFERS];

    ID3D12GraphicsCommandList   *command_list;
    ID3D12Fence                 *fence;
    HANDLE                       fence_event;
    UINT64                       fence_last_signaled;
};

static bool create_device(Renderer *r)
{
#ifdef DX12_ENABLE_DEBUG_LAYER
    ID3D12Debug *debug = nullptr;
    if (SUCCEEDED(D3D12GetDebugInterface(IID_PPV_ARGS(&debug)))) {
        debug->EnableDebugLayer();
        debug->Release();
    }
#endif

    if (D3D12CreateDevice(nullptr, D3D_FEATURE_LEVEL_11_0, IID_PPV_ARGS(&r->device)) != S_OK)
        return false;

#ifdef DX12_ENABLE_DEBUG_LAYER
    {
        ID3D12InfoQueue *info_queue = nullptr;
        if (SUCCEEDED(r->device->QueryInterface(IID_PPV_ARGS(&info_queue)))) {
            info_queue->SetBreakOnSeverity(D3D12_MESSAGE_SEVERITY_ERROR, true);
            info_queue->SetBreakOnSeverity(D3D12_MESSAGE_SEVERITY_CORRUPTION, true);
            info_queue->Release();
        }
    }
#endif

    return true;
}

static bool create_command_queue(Renderer *r)
{
    D3D12_COMMAND_QUEUE_DESC desc = {};
    desc.Type     = D3D12_COMMAND_LIST_TYPE_DIRECT;
    desc.Flags    = D3D12_COMMAND_QUEUE_FLAG_NONE;
    desc.NodeMask = 1;
    return r->device->CreateCommandQueue(&desc, IID_PPV_ARGS(&r->command_queue)) == S_OK;
}

static bool create_descriptor_heaps(Renderer *r)
{
    {
        D3D12_DESCRIPTOR_HEAP_DESC desc = {};
        desc.Type           = D3D12_DESCRIPTOR_HEAP_TYPE_RTV;
        desc.NumDescriptors = NUM_BACK_BUFFERS;
        desc.Flags          = D3D12_DESCRIPTOR_HEAP_FLAG_NONE;
        desc.NodeMask       = 1;
        if (r->device->CreateDescriptorHeap(&desc, IID_PPV_ARGS(&r->rtv_heap)) != S_OK)
            return false;

        SIZE_T rtv_size = r->device->GetDescriptorHandleIncrementSize(D3D12_DESCRIPTOR_HEAP_TYPE_RTV);
        D3D12_CPU_DESCRIPTOR_HANDLE handle = r->rtv_heap->GetCPUDescriptorHandleForHeapStart();
        for (int i = 0; i < NUM_BACK_BUFFERS; i++) {
            r->rtv_descriptors[i] = handle;
            handle.ptr += rtv_size;
        }
    }

    {
        D3D12_DESCRIPTOR_HEAP_DESC desc = {};
        desc.Type           = D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV;
        desc.NumDescriptors = SRV_HEAP_SIZE;
        desc.Flags          = D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE;
        if (r->device->CreateDescriptorHeap(&desc, IID_PPV_ARGS(&r->srv_heap)) != S_OK)
            return false;

        srv_alloc_init(&r->srv_alloc, r->device, r->srv_heap);
    }

    return true;
}

static bool create_frame_resources(Renderer *r)
{
    for (int i = 0; i < r->frame_count; i++) {
        if (r->device->CreateCommandAllocator(D3D12_COMMAND_LIST_TYPE_DIRECT,
                IID_PPV_ARGS(&r->frames[i].command_allocator)) != S_OK)
            return false;
        r->frames[i].fence_value = 0;
    }

    if (r->device->CreateCommandList(0, D3D12_COMMAND_LIST_TYPE_DIRECT,
            r->frames[0].command_allocator, nullptr, IID_PPV_ARGS(&r->command_list)) != S_OK)
        return false;
    if (r->command_list->Close() != S_OK)
        return false;

    if (r->device->CreateFence(0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&r->fence)) != S_OK)
        return false;

    r->fence_event = CreateEventW(nullptr, FALSE, FALSE, nullptr);
    return r->fence_event != nullptr;
}

static bool create_swap_chain(Renderer *r)
{
    DXGI_SWAP_CHAIN_DESC1 sd = {};
    sd.BufferCount        = NUM_BACK_BUFFERS;
    sd.Width              = 0;
    sd.Height             = 0;
    sd.Format             = DXGI_FORMAT_R8G8B8A8_UNORM;
    sd.Flags              = DXGI_SWAP_CHAIN_FLAG_FRAME_LATENCY_WAITABLE_OBJECT;
    sd.BufferUsage        = DXGI_USAGE_RENDER_TARGET_OUTPUT;
    sd.SampleDesc.Count   = 1;
    sd.SampleDesc.Quality = 0;
    sd.SwapEffect         = DXGI_SWAP_EFFECT_FLIP_DISCARD;
    sd.AlphaMode          = DXGI_ALPHA_MODE_UNSPECIFIED;
    sd.Scaling            = DXGI_SCALING_STRETCH;
    sd.Stereo             = FALSE;

    IDXGIFactory5 *factory = nullptr;
    if (CreateDXGIFactory1(IID_PPV_ARGS(&factory)) != S_OK)
        return false;

    BOOL allow_tearing = FALSE;
    factory->CheckFeatureSupport(DXGI_FEATURE_PRESENT_ALLOW_TEARING, &allow_tearing, sizeof(allow_tearing));
    r->tearing_support = (allow_tearing == TRUE);
    if (r->tearing_support)
        sd.Flags |= DXGI_SWAP_CHAIN_FLAG_ALLOW_TEARING;

    IDXGISwapChain1 *swap_chain1 = nullptr;
    if (factory->CreateSwapChainForHwnd(r->command_queue, r->hwnd, &sd, nullptr, nullptr, &swap_chain1) != S_OK) {
        factory->Release();
        return false;
    }

    if (swap_chain1->QueryInterface(IID_PPV_ARGS(&r->swap_chain)) != S_OK) {
        swap_chain1->Release();
        factory->Release();
        return false;
    }

    if (r->tearing_support)
        factory->MakeWindowAssociation(r->hwnd, DXGI_MWA_NO_ALT_ENTER);

    swap_chain1->Release();
    factory->Release();

    r->swap_chain->SetMaximumFrameLatency(NUM_BACK_BUFFERS);
    r->swap_chain_waitable = r->swap_chain->GetFrameLatencyWaitableObject();
    return true;
}

static void create_render_targets(Renderer *r)
{
    for (UINT i = 0; i < NUM_BACK_BUFFERS; i++) {
        ID3D12Resource *back_buffer = nullptr;
        r->swap_chain->GetBuffer(i, IID_PPV_ARGS(&back_buffer));
        r->device->CreateRenderTargetView(back_buffer, nullptr, r->rtv_descriptors[i]);
        r->render_targets[i] = back_buffer;
    }
}

static void cleanup_render_targets(Renderer *r)
{
    for (UINT i = 0; i < NUM_BACK_BUFFERS; i++) {
        if (r->render_targets[i]) {
            r->render_targets[i]->Release();
            r->render_targets[i] = nullptr;
        }
    }
}

static void wait_for_pending(Renderer *r)
{
    r->command_queue->Signal(r->fence, ++r->fence_last_signaled);
    r->fence->SetEventOnCompletion(r->fence_last_signaled, r->fence_event);
    WaitForSingleObject(r->fence_event, INFINITE);
}

static FrameContext *wait_for_next_frame(Renderer *r)
{
    FrameContext *fc = &r->frames[r->frame_index % r->frame_count];
    if (r->fence->GetCompletedValue() < fc->fence_value) {
        r->fence->SetEventOnCompletion(fc->fence_value, r->fence_event);
        HANDLE waitables[] = { r->swap_chain_waitable, r->fence_event };
        WaitForMultipleObjects(2, waitables, TRUE, INFINITE);
    } else {
        WaitForSingleObject(r->swap_chain_waitable, INFINITE);
    }
    return fc;
}

static void init_imgui(Renderer *r)
{
    IMGUI_CHECKVERSION();
    ImGui::CreateContext();
    ImGuiIO &io = ImGui::GetIO();
    io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard;
    io.ConfigFlags |= ImGuiConfigFlags_DockingEnable;

    ImGui::StyleColorsDark();

    ImGui_ImplWin32_Init(r->hwnd);

    ImGui_ImplDX12_InitInfo init_info = {};
    init_info.Device              = r->device;
    init_info.CommandQueue        = r->command_queue;
    init_info.NumFramesInFlight   = r->frame_count;
    init_info.RTVFormat           = DXGI_FORMAT_R8G8B8A8_UNORM;
    init_info.DSVFormat           = DXGI_FORMAT_UNKNOWN;
    init_info.SrvDescriptorHeap   = r->srv_heap;
    init_info.SrvDescriptorAllocFn = [](ImGui_ImplDX12_InitInfo *info, D3D12_CPU_DESCRIPTOR_HANDLE *out_cpu, D3D12_GPU_DESCRIPTOR_HANDLE *out_gpu) {
        Renderer *r = (Renderer *)info->UserData;
        srv_alloc_alloc(&r->srv_alloc, out_cpu, out_gpu);
    };
    init_info.SrvDescriptorFreeFn = [](ImGui_ImplDX12_InitInfo *info, D3D12_CPU_DESCRIPTOR_HANDLE cpu, D3D12_GPU_DESCRIPTOR_HANDLE gpu) {
        Renderer *r = (Renderer *)info->UserData;
        srv_alloc_free(&r->srv_alloc, cpu, gpu);
    };
    init_info.UserData = r;
    ImGui_ImplDX12_Init(&init_info);
}

Renderer *renderer_create(RendererDesc *desc)
{
    Renderer *r = new Renderer();
    memset(r, 0, sizeof(*r));

    r->hwnd        = (HWND)desc->window_handle;
    r->width       = desc->width;
    r->height      = desc->height;
    r->frame_count = desc->frame_count;
    if (r->frame_count > NUM_BACK_BUFFERS) r->frame_count = NUM_BACK_BUFFERS;

    if (!create_device(r))            goto fail;
    if (!create_command_queue(r))     goto fail;
    if (!create_descriptor_heaps(r))  goto fail;
    if (!create_frame_resources(r))   goto fail;
    if (!create_swap_chain(r))        goto fail;
    create_render_targets(r);

    init_imgui(r);
    return r;

fail:
    renderer_destroy(r);
    return nullptr;
}

void renderer_destroy(Renderer *r)
{
    if (!r) return;

    wait_for_pending(r);

    ImGui_ImplDX12_Shutdown();
    ImGui_ImplWin32_Shutdown();
    ImGui::DestroyContext();

    cleanup_render_targets(r);

    if (r->swap_chain)       { r->swap_chain->SetFullscreenState(false, nullptr); r->swap_chain->Release(); }
    if (r->swap_chain_waitable) CloseHandle(r->swap_chain_waitable);
    for (int i = 0; i < r->frame_count; i++)
        if (r->frames[i].command_allocator) r->frames[i].command_allocator->Release();
    if (r->command_queue)    r->command_queue->Release();
    if (r->command_list)     r->command_list->Release();
    if (r->rtv_heap)         r->rtv_heap->Release();
    if (r->srv_heap)         r->srv_heap->Release();
    if (r->fence)            r->fence->Release();
    if (r->fence_event)      CloseHandle(r->fence_event);
    if (r->device)           r->device->Release();

#ifdef DX12_ENABLE_DEBUG_LAYER
    IDXGIDebug1 *dxgi_debug = nullptr;
    if (SUCCEEDED(DXGIGetDebugInterface1(0, IID_PPV_ARGS(&dxgi_debug)))) {
        dxgi_debug->ReportLiveObjects(DXGI_DEBUG_ALL, DXGI_DEBUG_RLO_SUMMARY);
        dxgi_debug->Release();
    }
#endif

    delete r;
}

bool renderer_begin_frame(Renderer *r)
{
    if ((r->swap_chain_occluded && r->swap_chain->Present(0, DXGI_PRESENT_TEST) == DXGI_STATUS_OCCLUDED)
        || IsIconic(r->hwnd))
    {
        Sleep(10);
        return false;
    }
    r->swap_chain_occluded = false;

    ImGui_ImplDX12_NewFrame();
    ImGui_ImplWin32_NewFrame();
    ImGui::NewFrame();
    return true;
}

void renderer_end_frame(Renderer *r)
{
    ImGui::Render();

    FrameContext *fc = wait_for_next_frame(r);
    UINT back_buffer_idx = r->swap_chain->GetCurrentBackBufferIndex();

    fc->command_allocator->Reset();
    r->command_list->Reset(fc->command_allocator, nullptr);

    D3D12_RESOURCE_BARRIER barrier = {};
    barrier.Type                   = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
    barrier.Flags                  = D3D12_RESOURCE_BARRIER_FLAG_NONE;
    barrier.Transition.pResource   = r->render_targets[back_buffer_idx];
    barrier.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
    barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_PRESENT;
    barrier.Transition.StateAfter  = D3D12_RESOURCE_STATE_RENDER_TARGET;
    r->command_list->ResourceBarrier(1, &barrier);

    const float clear_color[4] = { 0.1f, 0.1f, 0.1f, 1.0f };
    r->command_list->ClearRenderTargetView(r->rtv_descriptors[back_buffer_idx], clear_color, 0, nullptr);
    r->command_list->OMSetRenderTargets(1, &r->rtv_descriptors[back_buffer_idx], FALSE, nullptr);
    r->command_list->SetDescriptorHeaps(1, &r->srv_heap);

    ImGui_ImplDX12_RenderDrawData(ImGui::GetDrawData(), r->command_list);

    barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_RENDER_TARGET;
    barrier.Transition.StateAfter  = D3D12_RESOURCE_STATE_PRESENT;
    r->command_list->ResourceBarrier(1, &barrier);
    r->command_list->Close();

    r->command_queue->ExecuteCommandLists(1, (ID3D12CommandList *const *)&r->command_list);
    r->command_queue->Signal(r->fence, ++r->fence_last_signaled);
    fc->fence_value = r->fence_last_signaled;

    HRESULT hr = r->swap_chain->Present(1, 0);
    r->swap_chain_occluded = (hr == DXGI_STATUS_OCCLUDED);
    r->frame_index++;
}

void renderer_resize(Renderer *r, int32_t width, int32_t height)
{
    if (width <= 0 || height <= 0) return;

    wait_for_pending(r);
    cleanup_render_targets(r);

    DXGI_SWAP_CHAIN_DESC1 desc = {};
    r->swap_chain->GetDesc1(&desc);
    r->swap_chain->ResizeBuffers(0, (UINT)width, (UINT)height, desc.Format, desc.Flags);

    create_render_targets(r);

    r->width  = width;
    r->height = height;
}
