#include "render_mesh.h"

#include "runtime/config.h"
#include "runtime/ds.h"
#include "runtime/mem_arena.h"
#include "runtime/threading.h"

#include <stdlib.h>
#include <string.h>

RT_CVAR_SZ(rt_MeshPoolSlots, "Number of slots available for meshes. Default: 1024", 1024);
RT_CVAR_SZ(rt_MeshletPoolSlots, "Number of slots available for meshlets. Default: 4096", 4096);
RT_CVAR_SZ(rt_MeshletStreamQueueSize, "Number of slots in the streaming queue. Default: 256", 256);
RT_CVAR_F(rt_MeshletStreamTimeslice,
          "Timeslice available for streaming meshlets. Default: 4ms",
          0.04f);
RT_CVAR_I(rt_MeshletDirectReleaseCount,
          "Number of meshlets that will get released immediately. Default: 1",
          1);

typedef struct rt_mesh_container {
    _Alignas(4) uint32_t refcount;
    struct rt_mesh_container *next_free;
    rt_render_mesh mesh;
} rt_mesh_container;

typedef struct rt_meshlet_container {
    struct rt_meshlet_container *next_free;
    rt_render_meshlet meshlet;
} rt_meshlet_container;

typedef struct {
    rt_meshlet_container *meshlet;
} rt_meshlet_release_queue_entry;

typedef struct {
    rt_resource_id meshlet_resource;
    rt_meshlet_container *destination;
} rt_meshlet_stream_queue_entry;

static unsigned int _current_frame_id = 0;
static rt_mesh_container *_mesh_container;
static rt_mesh_container *_first_free_mesh;
static rt_meshlet_container *_meshlet_container;
static rt_meshlet_container *_first_free_meshlet;

static rt_minheap _reclaim_heap;
static void *_reclaim_heap_memory;
static rt_hashtable _mesh_lut;
static void *_mesh_lut_memory;
static rt_mutex *_mesh_lock;

static rt_queue _stream_queue;
static void *_stream_queue_memory;
static rt_queue _release_queue;
static void *_release_queue_memory;
static rt_mutex *_meshlet_lock;

extern rt_result InitMeshletPools(void);
extern void ShutdownMeshletPools(void);

rt_result rtInitMeshStreaming(void) {
    rtRegisterCVAR(&rt_MeshPoolSlots);
    rtRegisterCVAR(&rt_MeshletPoolSlots);
    rtRegisterCVAR(&rt_MeshletStreamQueueSize);
    rtRegisterCVAR(&rt_MeshletStreamTimeslice);
    rtRegisterCVAR(&rt_MeshletDirectReleaseCount);

    if (!RT_IS_POWER_OF_TWO(rt_MeshPoolSlots.sz)) {
        rtReportError("RENCOM", "rt_MeshPoolSlots must be a power of two.");
        return RT_INVALID_VALUE;
    }

    if (rt_MeshletDirectReleaseCount.i == 0) {
        rtLog("RENCOM", "rt_MeshletDirectReleaseCount must at least be one.");
        rt_MeshletDirectReleaseCount.i = 1;
    }

    _mesh_container = calloc(rt_MeshPoolSlots.sz, sizeof(rt_mesh_container));
    if (!_mesh_container)
        return RT_OUT_OF_MEMORY;
    for (size_t i = 0; i < rt_MeshPoolSlots.sz - 1; ++i) {
        _mesh_container[i].next_free = &_mesh_container[i + 1];
    }
    _first_free_mesh = _mesh_container;

    _reclaim_heap_memory = malloc((sizeof(int) + sizeof(size_t)) * rt_MeshPoolSlots.sz);
    if (!_reclaim_heap_memory) {
        free(_mesh_container);
        return RT_OUT_OF_MEMORY;
    }
    _reclaim_heap = rtCreateMinheap((int *)_reclaim_heap_memory,
                                    (int *)+rt_MeshPoolSlots.sz,
                                    sizeof(size_t),
                                    rt_MeshPoolSlots.sz,
                                    0);

    _mesh_lut_memory = malloc(RT_HASH_TABLE_MEMORY_REQUIRED(2 * rt_MeshPoolSlots.sz));
    if (!_mesh_lut_memory) {
        free(_reclaim_heap_memory);
        free(_mesh_container);
        return RT_OUT_OF_MEMORY;
    }
    _mesh_lut = rtCreateHashtable(2 * rt_MeshPoolSlots.sz, _mesh_lut_memory, NULL, NULL);

    _meshlet_container = calloc(rt_MeshletPoolSlots.sz, sizeof(rt_meshlet_container));
    if (!_meshlet_container) {
        free(_reclaim_heap_memory);
        free(_mesh_container);
        return RT_OUT_OF_MEMORY;
    }
    for (size_t i = 0; i < rt_MeshPoolSlots.sz - 1; ++i) {
        _meshlet_container[i].next_free = &_meshlet_container[i + 1];
    }
    _first_free_meshlet = _meshlet_container;

    _stream_queue_memory =
        calloc(rt_MeshletStreamQueueSize.sz, sizeof(rt_meshlet_stream_queue_entry));
    if (!_stream_queue_memory) {
        free(_reclaim_heap_memory);
        free(_meshlet_container);
        free(_mesh_container);
        return RT_OUT_OF_MEMORY;
    }
    _stream_queue = rtCreateQueue(rt_MeshletStreamQueueSize.sz,
                                  _stream_queue_memory,
                                  sizeof(rt_meshlet_stream_queue_entry));

    _release_queue_memory =
        calloc(rt_MeshletStreamQueueSize.sz, sizeof(rt_meshlet_release_queue_entry));
    if (!_release_queue_memory) {
        free(_meshlet_container);
        free(_mesh_container);
        free(_stream_queue_memory);
        return RT_OUT_OF_MEMORY;
    }
    _release_queue = rtCreateQueue(rt_MeshletStreamQueueSize.sz,
                                   _release_queue_memory,
                                   sizeof(rt_meshlet_release_queue_entry));

    _mesh_lock    = rtCreateMutex();
    _meshlet_lock = rtCreateMutex();

    rt_result res = InitMeshletPools();
    if (res != RT_SUCCESS) {
        rtShutdownMeshStreaming();
        return res;
    }

    return RT_SUCCESS;
}

void rtShutdownMeshStreaming(void) {
    free(_mesh_container);
    free(_meshlet_container);
    free(_stream_queue_memory);
    free(_release_queue_memory);
    free(_reclaim_heap_memory);
    free(_mesh_lut_memory);

    ShutdownMeshletPools();
}

static rt_meshlet_container *AllocMeshletContainer(void) {
    rtLockMutex(_meshlet_lock);
    rt_meshlet_container *container = _first_free_meshlet;
    if (container) {
        container->meshlet.state = RT_RENDER_MESHLET_STATE_LOADING;
        _first_free_meshlet      = container->next_free;
    }
    rtUnlockMutex(_meshlet_lock);
    return container;
}

static rt_result LoadMeshlet(rt_resource_id meshlet_id, rt_meshlet_container *destination) {
    size_t resource_size = rtGetResourceSize(meshlet_id);
    if (!resource_size)
        return RT_INVALID_VALUE;

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

    void *resource_buffer = rtArenaPush(temp.arena, resource_size);
    if (!resource_buffer) {
        rtLog("RENCOM",
              "Unable to allocate %zu bytes for loading meshlet %llx",
              resource_size,
              meshlet_id);
        rtReturnTemporaryArena(temp);
        return RT_OUT_OF_MEMORY;
    }

    rt_result res = rtGetResource(meshlet_id, resource_buffer);
    if (res != RT_SUCCESS) {
        rtLog("RENCOM", "Failed to get meshlet resource %llx", meshlet_id);
    }
    rt_resource *meshlet_resource = resource_buffer;
    if (!RT_VERIFY(meshlet_resource->type == RT_RESOURCE_MESHLET)) {
        rtLog("RENCOM", "Resource %llx is not a meshlet.", meshlet_id);
        res = RT_INVALID_VALUE;
        goto out;
    }

out:
    rtReturnTemporaryArena(temp);
    return res;
}

static void ReleaseMeshlet(rt_meshlet_container *container) {
    /* TODO: Release vertex and index buffers */
    container->meshlet.state = RT_RENDER_MESHLET_STATE_INVALID;

    rtLockMutex(_meshlet_lock);
    container->next_free = _first_free_meshlet;
    _first_free_meshlet  = container;
    rtUnlockMutex(_meshlet_lock);
}

static int RecursiveReleaseMeshlets(rt_render_meshlet *meshlet, int release_count) {
    for (unsigned int i = 0; i < meshlet->static_data.child_count; ++i) {
        release_count = RecursiveReleaseMeshlets(meshlet->children[i], release_count);
    }

    rt_meshlet_container *container =
        (rt_meshlet_container *)((char *)meshlet - offsetof(rt_meshlet_container, meshlet));
    rtLockMutex(_meshlet_lock);
    if (release_count >= rt_MeshletDirectReleaseCount.i && !rtQueueIsFull(&_release_queue)) {
        /* Enqueue for later */
        rt_meshlet_release_queue_entry entry = {.meshlet = container};
        rtQueuePush(&_release_queue, &entry);
        rtUnlockMutex(_meshlet_lock);
        return release_count;
    } else {
        rtUnlockMutex(_meshlet_lock);
        ReleaseMeshlet(container);
        return release_count + 1;
    }
}

rt_render_mesh *rtGetRenderMesh(rt_resource_id mesh_id) {
    size_t mesh_resource_size = rtGetResourceSize(mesh_id);
    if (!mesh_resource_size) {
        rtLog("RENCOM", "Tried to load invalid mesh %llx");
        return NULL;
    }

    rtLockMutex(_mesh_lock);
    /* Check if this mesh is already loaded */
    size_t mesh_index = 0;
    if ((mesh_index = rtHashtableLookup(&_mesh_lut, mesh_id, SIZE_MAX)) != SIZE_MAX) {
        rt_mesh_container *container = &_mesh_container[mesh_id];
        container->refcount++;
        rtUnlockMutex(_mesh_lock);
        return &container->mesh;
    }

    if (!_first_free_mesh) {
        /* Try to reclaim used space */
        if (rtMinheapIsEmpty(&_reclaim_heap)) {
            rtLog("RENCOM", "Tried to load new mesh %llx, but no space is available.", mesh_id);
            rtUnlockMutex(_mesh_lock);
            return NULL;
        }

        size_t reuse_idx;
        rtMinheapPop(&_reclaim_heap, &reuse_idx);
        RT_ASSERT(_mesh_container[reuse_idx].refcount == 0,
                  "In-use mesh container found in reclaim heap.");
        _mesh_container[reuse_idx].next_free = NULL;
        _first_free_mesh                     = &_mesh_container[reuse_idx];

        /* Free all associated meshlets */
        RT_VERIFY(_mesh_container[reuse_idx].mesh.root);
        int num_released = RecursiveReleaseMeshlets(_mesh_container[reuse_idx].mesh.root, 0);
        RT_ASSERT(num_released >= 1, "At least one meshlet should have been released.");
    }

    rt_mesh_container *container = _first_free_mesh;
    _first_free_mesh             = container->next_free;

    mesh_index = container - _mesh_container;
    rtHashtableInsert(&_mesh_lut, mesh_id, mesh_index);

    rtUnlockMutex(_mesh_lock);

    container->refcount  = 1;
    container->next_free = NULL;

    rt_temp_arena temp = rtGetTemporaryArena(NULL, 0);
    if (!temp.arena)
        return NULL;
    void *resource_buffer = rtArenaPush(temp.arena, mesh_resource_size);
    if (!resource_buffer) {
        rtLog("RENCOM",
              "Tried to load new mesh %llx, but no memory for storing the mesh resource is "
              "available.",
              mesh_id);
        rtLockMutex(_mesh_lock);
        container->refcount  = 0;
        container->next_free = _first_free_mesh;
        _first_free_mesh     = container;
        rtHashtableRemove(&_mesh_lut, mesh_id);
        rtUnlockMutex(_mesh_lock);
        return NULL;
    }

    if (rtGetResource(mesh_id, resource_buffer) != RT_SUCCESS) {
        rtLog("RENCOM", "Failed to load mesh %llx", mesh_id);
        rtUnlockMutex(_mesh_lock);
        return NULL;
    }
    rt_resource *resource = resource_buffer;
    RT_VERIFY(resource->type == RT_RESOURCE_MESH);
    memcpy(&container->mesh.static_data, resource->data, sizeof(rt_render_mesh_static_data));
    rtReturnTemporaryArena(temp);
    resource        = NULL;
    resource_buffer = NULL;

    /* Enqueue the root load */
    rt_meshlet_stream_queue_entry root_load = {0};
    root_load.destination                   = AllocMeshletContainer();
    root_load.meshlet_resource              = container->mesh.static_data.root_resource;
    if (!root_load.destination) {
        /* Oh no, we can't load this */
        rtLog(
            "RENCOM",
            "Tried to load new mesh %llx, but no space for loading the root meshlet is available.",
            mesh_id);
        rtLockMutex(_mesh_lock);
        container->refcount  = 0;
        container->next_free = _first_free_mesh;
        _first_free_mesh     = container;
        rtHashtableRemove(&_mesh_lut, mesh_id);
        rtUnlockMutex(_mesh_lock);
        return NULL;
    }
    container->mesh.root = &root_load.destination->meshlet;
    rtQueuePush(&_stream_queue, &root_load);

    return &container->mesh;
}

void rtReleaseRenderMesh(rt_resource_id mesh_id) {
    rtLockMutex(_mesh_lock);
    size_t mesh_index;
    if ((mesh_index = rtHashtableLookup(&_mesh_lut, mesh_id, SIZE_MAX)) == SIZE_MAX) {
        rtUnlockMutex(_mesh_lock);
        return;
    }
    rt_mesh_container *container = &_mesh_container[mesh_index];
    if ((--container->refcount) == 0) {
        RT_VERIFY(container->mesh.root);
        RecursiveReleaseMeshlets(container->mesh.root, 0);
        container->next_free = _first_free_mesh;
        _first_free_mesh     = container;

        rtMinheapPush(&_reclaim_heap, (int)_current_frame_id, &mesh_index);
    }
    rtUnlockMutex(_mesh_lock);
}

void rtMeshStreamingNewFrame(unsigned int frame_id) {
    _current_frame_id = frame_id;

    float time = 0.f;
    while (time < rt_MeshletStreamTimeslice.f) {
        rt_meshlet_release_queue_entry release = {NULL};
        rt_meshlet_stream_queue_entry stream   = {0};
        bool have_release = false, have_stream = false;

        rtLockMutex(_meshlet_lock);
        if (!rtQueueIsEmpty(&_release_queue)) {
            rtQueuePop(&_release_queue, &release);
            have_release = true;
        }
        if (!rtQueueIsEmpty(&_stream_queue)) {
            rtQueuePop(&_stream_queue, &stream);
            have_stream = true;
        }
        rtUnlockMutex(_meshlet_lock);
        if (!have_stream && !have_release)
            break;

        if (have_release) {
            ReleaseMeshlet(release.meshlet);
        }
        if (have_stream) {
            if (LoadMeshlet(stream.meshlet_resource, stream.destination) != RT_SUCCESS) {
                stream.destination->meshlet.state = RT_RENDER_MESHLET_STATE_LOAD_FAILED;
            }
        }
    }
}
