#include "runtime/config.h"
#include "runtime/handles.h"
#include "runtime/mem_arena.h"
#include "runtime/renderer_api.h"
#include "runtime/resources.h"
#include "runtime/threading.h"

#include "gpu.h"
#include "pipelines.h"

#include <stdlib.h>
#include <volk/volk.h>

RT_CVAR_I(r_VkMaxPipelineCount, "Maximum number of pipeline objects. Default: 1024", 1024);

typedef struct rt_pipeline_s {
    uint32_t version;
    rt_pipeline pipeline;
    struct rt_pipeline_s *next_free;
} rt_pipeline_slot;

static rt_pipeline_slot *_pipelines;
static rt_pipeline_slot *_first_free;
static rt_rwlock _lock;

static void DestroyPipeline(rt_pipeline_slot *slot) {
    if (slot->pipeline.pipeline) {
        vkDestroyPipeline(g_gpu.device, slot->pipeline.pipeline, g_gpu.alloc_cb);
    }
    slot->next_free = _first_free;
    _first_free     = slot;
}

static VkShaderModule CreateShaderModuleFromResource(rt_resource_id rid) {
    if (rid == RT_INVALID_RESOURCE_ID)
        return VK_NULL_HANDLE;
    rt_resource *resource = NULL;
    size_t size           = rtGetResourceSize(rid);
    if (size == 0)
        return VK_NULL_HANDLE;

    rt_temp_arena temp = rtGetTemporaryArena(NULL, 0);
    if (!temp.arena)
        return VK_NULL_HANDLE;

    VkShaderModule module = VK_NULL_HANDLE;
    resource = rtArenaPush(temp.arena, size);
    if (!resource) {
        rtLog("VK", "Failed to allocate temporary memory for retrieving a shader resource");
        goto out;
    }

    if (rtGetResource(rid, resource) != RT_SUCCESS) {
        goto out;
    }
    if (resource->type != RT_RESOURCE_SHADER) {
        rtLog("VK", "Attempted to create a shader module from a non-shader resource %llx", rid);
        goto out;
    }
    rt_shader_info *info = resource->data;
    if (!info) {
        rtLog("VK", "Shader resource %llx has no attached shader_info", rid);
        goto out;
    }

    VkShaderModuleCreateInfo module_info = {.sType    = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
                                            .pCode    = rtResolveRelptr(&info->bytecode),
                                            .codeSize = info->bytecode_length};
    if (vkCreateShaderModule(g_gpu.device, &module_info, g_gpu.alloc_cb, &module) != VK_SUCCESS) {
        rtLog("VK", "Failed to create the shader module from resource %llx", rid);
        goto out;
    }
out:
    rtReturnTemporaryArena(temp);
    return module;
}

static bool CreateComputePipeline(VkShaderModule compute_shader,
                                  const rt_pipeline_info *info,
                                  rt_pipeline_slot *slot) {
    return false;
}

static bool CreateGraphicsPipeline(VkShaderModule vertex_shader,
                                   VkShaderModule fragment_shader,
                                   const rt_pipeline_info *info,
                                   rt_pipeline_slot *slot) {
    return false;
}

rt_result InitPipelineManagement(void) {
    rt_create_rwlock_result lock_res = rtCreateRWLock();
    if (!lock_res.ok)
        return RT_UNKNOWN_ERROR;
    _lock = lock_res.lock;

    _pipelines = calloc(r_VkMaxPipelineCount.i, sizeof(rt_pipeline_slot));
    if (!_pipelines) {
        rtDestroyRWLock(&_lock);
        return RT_OUT_OF_MEMORY;
    }
    /* Keep [0] unused to preserve 0 as the invalid handle */
    _first_free = &_pipelines[1];
    for (int i = 1; i < r_VkMaxPipelineCount.i - 1; ++i) {
        _pipelines[i].next_free = &_pipelines[i + 1];
    }
    return RT_SUCCESS;
}

void ShutdownPipelineManagement(void) {
    for (int i = 1; i < r_VkMaxPipelineCount.i; ++i) {
        DestroyPipeline(&_pipelines[i]);
    }
    free(_pipelines);
    rtDestroyRWLock(&_lock);
    _first_free = NULL;
}

rt_pipeline_handle RT_RENDERER_API_FN(CompilePipeline)(const rt_pipeline_info *info) {
    rt_pipeline_handle handle = RT_INVALID_HANDLE;
    rtLockWrite(&_lock);
    if (!_first_free) {
        rtLog("VK", "No free pipeline slots!");
        rtUnlockWrite(&_lock);
        return handle;
    }
    rt_pipeline_slot *slot = _first_free;
    _first_free            = slot->next_free;
    slot->version          = (slot->version + 1) & RT_GFX_HANDLE_MAX_VERSION;

    /* No other thread that calls compile gets the same slot.
     * Another thread accessing the slot via GetPipeline would get a version mismatch.
     * The same holds for DestroyPipeline
     */
    rtUnlockWrite(&_lock);

    VkShaderModule vertex_shader   = CreateShaderModuleFromResource(info->vertex_shader);
    VkShaderModule fragment_shader = CreateShaderModuleFromResource(info->fragment_shader);
    VkShaderModule compute_shader  = CreateShaderModuleFromResource(info->compute_shader);

    RT_UNUSED(vertex_shader);
    RT_UNUSED(fragment_shader);
    RT_UNUSED(compute_shader);

    bool create_success = false;
    if (compute_shader) {
        create_success = CreateComputePipeline(compute_shader, info, slot);
    } else if (vertex_shader && fragment_shader) {
        create_success = CreateGraphicsPipeline(vertex_shader, fragment_shader, info, slot);
    } else {
        rtLog("VK", "Invalid combination of shaders in pipeline info.");
    }

    if (create_success) {
        handle.version = slot->version;
        handle.index   = (uint32_t)(slot - _pipelines);
    }

    return handle;
}

void RT_RENDERER_API_FN(DestroyPipeline)(rt_pipeline_handle handle) {
    if (handle.index >= (uint32_t)r_VkMaxPipelineCount.i)
        return;
    rtLockWrite(&_lock);
    if (_pipelines[handle.index].version == handle.version)
        DestroyPipeline(&_pipelines[handle.index]);
    else
        rtLog("VK", "Tried to destroy a pipeline using an outdated handle.");
    rtUnlockWrite(&_lock);
}

const rt_pipeline *rtGetPipeline(rt_pipeline_handle handle) {
    if (handle.index >= (uint32_t)r_VkMaxPipelineCount.i)
        return NULL;
    rtLockRead(&_lock);
    rt_pipeline *res = NULL;
    if (_pipelines[handle.index].version == handle.version)
        res = &_pipelines[handle.index].pipeline;
    else
        rtLog("VK", "Tried to access a pipeline using an outdated handle.");
    rtUnlockRead(&_lock);
    return res;
}