#include <d3d11.h>
#include <d3d11_1.h>

#include "gfx/effect.h"
#include "gfx/renderer_api.h"
#include "runtime/config.h"
#include "runtime/handles.h"
#include "runtime/mem_arena.h"
#include "runtime/threading_helpers.hpp"

#include "device_objects.hpp"
#include "gpu.hpp"

RT_CVAR_I(rt_Dx11MaxPipelines,
          "Maximum number of simultaneously existing pipelines. Default: 128",
          128);

static rt_pipeline *_pipelines;
static rt_pipeline *_first_free;
static rt_mutex *_lock;

rt_result InitPipelineManagement() {
    _pipelines =
        reinterpret_cast<rt_pipeline *>(calloc((size_t)rt_Dx11MaxPipelines.i, sizeof(rt_pipeline)));
    if (!_pipelines)
        return RT_OUT_OF_MEMORY;

    _first_free = _pipelines + 1;
    for (int i = 0; i < rt_Dx11MaxPipelines.i - 1; ++i)
        _pipelines[i].next_free = &_pipelines[i + 1];

    _lock = rtCreateMutex();
    if (!_lock) {
        free(_pipelines);
        return RT_UNKNOWN_ERROR;
    }

    return RT_SUCCESS;
}

void ShutdownPipelineManagement() {
    for (int i = 0; i < rt_Dx11MaxPipelines.i; ++i) {
        if (_pipelines[i].compute_shader)
            _pipelines[i].compute_shader->Release();
        if (_pipelines[i].vertex_shader)
            _pipelines[i].vertex_shader->Release();
        if (_pipelines[i].pixel_shader)
            _pipelines[i].pixel_shader->Release();
        if (_pipelines[i].input_layout)
            _pipelines[i].input_layout->Release();
    }
    free(_pipelines);
    rtDestroyMutex(_lock);
}

rt_result GetShader(rt_resource_id id, rt_shader_info **p_shader, rt_arena *arena) {
    size_t shader_size = rtGetResourceSize(id);
    if (shader_size == 0)
        return RT_INVALID_VALUE;
    void *buffer = rtArenaPush(arena, shader_size);
    if (!buffer)
        return RT_OUT_OF_MEMORY;
    rt_result res = rtGetResource(id, buffer);
    if (res != RT_SUCCESS) {
        rtArenaPop(arena, shader_size);
        return res;
    }

    rt_resource *resource = reinterpret_cast<rt_resource *>(buffer);
    RT_ASSERT(resource->type == RT_RESOURCE_SHADER, "Expected a shader");
    *p_shader = reinterpret_cast<rt_shader_info *>(resource->data);
    return RT_SUCCESS;
}

extern "C" rt_pipeline_handle RT_RENDERER_API_FN(CompilePipeline)(const rt_pipeline_info *info) {
    rt_pipeline *slot = nullptr;
    {
        auto lg = rtAutoLock(_lock);

        slot = _first_free;
        if (slot)
            _first_free = slot->next_free;
    }
    if (!slot) {
        rtLog("dx11", "Could not create pipeline, because no slots are available.");
        return RT_INVALID_HANDLE;
    }

    rt_temp_arena temp = rtGetTemporaryArena(NULL, 0);

    if (info->vertex_shader != RT_INVALID_RESOURCE_ID) {
        rt_shader_info *vs;
        if (GetShader(info->vertex_shader, &vs, temp.arena) != RT_SUCCESS) {
            rtReportError("dx11", "Could not retrieve vertex shader data.");
            auto lg         = rtAutoLock(_lock);
            slot->next_free = _first_free;
            _first_free     = slot;
            return RT_INVALID_HANDLE;
        }

        void *bytecode = rtResolveRelptr(&vs->bytecode);
        if (FAILED(g_gpu.device->CreateVertexShader(bytecode,
                                                    vs->bytecode_length,
                                                    NULL,
                                                    &slot->vertex_shader))) {
            rtReportError("dx11", "Vertex shader creation failed");
            auto lg         = rtAutoLock(_lock);
            slot->next_free = _first_free;
            _first_free     = slot;
            return RT_INVALID_HANDLE;
        }

        // TODO: effects should specify the expected vertex layout
        // For now, we use a default
        /* clang-format off */
        D3D11_INPUT_ELEMENT_DESC default_layout[] = {
            {"POSITION", 0,  DXGI_FORMAT_R32G32B32_FLOAT, 0,                            0, D3D11_INPUT_PER_VERTEX_DATA, 0},
            {"NORMAL",   0,  DXGI_FORMAT_R32G32B32_FLOAT, 0, D3D11_APPEND_ALIGNED_ELEMENT, D3D11_INPUT_PER_VERTEX_DATA, 0},
            {"TANGENT",  0,  DXGI_FORMAT_R32G32B32_FLOAT, 0, D3D11_APPEND_ALIGNED_ELEMENT, D3D11_INPUT_PER_VERTEX_DATA, 0},
            {"TEXCOORD", 0, DXGI_FORMAT_R32G32_FLOAT,     0, D3D11_APPEND_ALIGNED_ELEMENT, D3D11_INPUT_PER_VERTEX_DATA, 0},
        };
        /* clang-format on */

        if (FAILED(g_gpu.device->CreateInputLayout(default_layout,
                                                   RT_ARRAY_COUNT(default_layout),
                                                   bytecode,
                                                   vs->bytecode_length,
                                                   &slot->input_layout))) {
            rtReportError("dx11", "Failed to create the vertex layout.");
            auto lg         = rtAutoLock(_lock);
            slot->next_free = _first_free;
            _first_free     = slot;
            return RT_INVALID_HANDLE;
        }
    }

    if (info->fragment_shader != RT_INVALID_RESOURCE_ID) {
        rt_shader_info *vs;
        if (GetShader(info->fragment_shader, &vs, temp.arena) != RT_SUCCESS) {
            rtReportError("dx11", "Could not retrieve fragment shader data.");
            auto lg         = rtAutoLock(_lock);
            slot->next_free = _first_free;
            _first_free     = slot;
            return RT_INVALID_HANDLE;
        }

        void *bytecode = rtResolveRelptr(&vs->bytecode);
        if (FAILED(g_gpu.device->CreatePixelShader(bytecode,
                                                   vs->bytecode_length,
                                                   NULL,
                                                   &slot->pixel_shader))) {
            rtReportError("dx11", "Fragment shader creation failed");
            auto lg         = rtAutoLock(_lock);
            slot->next_free = _first_free;
            _first_free     = slot;
            return RT_INVALID_HANDLE;
        }
    }

    if (info->compute_shader != RT_INVALID_RESOURCE_ID) {
        rt_shader_info *vs;
        if (GetShader(info->compute_shader, &vs, temp.arena) != RT_SUCCESS) {
            rtReportError("dx11", "Could not retrieve compute shader data.");
            auto lg         = rtAutoLock(_lock);
            slot->next_free = _first_free;
            _first_free     = slot;
            return RT_INVALID_HANDLE;
        }

        void *bytecode = rtResolveRelptr(&vs->bytecode);
        if (FAILED(g_gpu.device->CreateComputeShader(bytecode,
                                                     vs->bytecode_length,
                                                     NULL,
                                                     &slot->compute_shader))) {
            rtReportError("dx11", "Compute shader creation failed");
            auto lg         = rtAutoLock(_lock);
            slot->next_free = _first_free;
            _first_free     = slot;
            return RT_INVALID_HANDLE;
        }
    }

    slot->version  = (slot->version + 1) % RT_RENDER_BACKEND_HANDLE_MAX_VERSION;
    uint32_t index = static_cast<uint32_t>(slot - _pipelines);
    return {slot->version, index};
}

extern "C" void RT_RENDERER_API_FN(DestroyPipeline)(rt_pipeline_handle handle) {
    if (!RT_IS_HANDLE_VALID(handle) || (int)handle.index >= rt_Dx11MaxPipelines.i)
        return;
    auto lg = rtAutoLock(_lock);
    if (handle.version != _pipelines[handle.index].version)
        return;

    if (_pipelines[handle.index].compute_shader)
        _pipelines[handle.index].compute_shader->Release();
    if (_pipelines[handle.index].vertex_shader)
        _pipelines[handle.index].vertex_shader->Release();
    if (_pipelines[handle.index].pixel_shader)
        _pipelines[handle.index].pixel_shader->Release();
    if (_pipelines[handle.index].input_layout)
        _pipelines[handle.index].input_layout->Release();

    _pipelines[handle.index].next_free = _first_free;
    _pipelines[handle.index].version =
        (_pipelines[handle.index].version + 1) % RT_RENDER_BACKEND_HANDLE_MAX_VERSION;
    _first_free = &_pipelines[handle.index];
}

rt_pipeline *rtGetPipeline(rt_pipeline_handle handle) {
    if (!RT_IS_HANDLE_VALID(handle) || (int)handle.index >= rt_Dx11MaxPipelines.i)
        return nullptr;
    auto lg = rtAutoLock(_lock);
    if (handle.version != _pipelines[handle.index].version)
        return nullptr;
    return &_pipelines[handle.index];
}