#include <meshoptimizer.h>
#pragma warning(push)
#pragma warning(disable : 4530)
#define TINYOBJLOADER_IMPLEMENTATION
#include <tiny_obj_loader.h>
#pragma warning(pop)

#include "meshlet_generator.hpp"
#include <runtime/mem_arena.h>

meshlet_generator::meshlet_generator()
    : m_meshes(nullptr), m_num_meshes(0u), m_meshlets(nullptr), m_num_meshlets(0u) {
}

meshlet_generator::~meshlet_generator() {
    Release();
}

void meshlet_generator::Release() {
    for (uint32_t i = 0; i < m_num_meshes; ++i) {
        delete m_meshes[i].vertices;
        delete m_meshes[i].indices;
    }
    delete m_meshes;
    m_meshes     = nullptr;
    m_num_meshes = 0;

    for (uint32_t i = 0; i < m_num_meshlets; ++i) {
        delete m_meshlets[i].indices;
        delete m_meshlets[i].vertices;
    }
    delete m_meshlets;
    m_meshlets     = nullptr;
    m_num_meshlets = 0;
}

rt_result meshlet_generator::LoadObj(const char *path) {
    tinyobj::ObjReaderConfig config;
    config.mtl_search_path = "./";
    config.triangulate     = true;
    tinyobj::ObjReader reader;

    if (!reader.ParseFromFile(path, config)) {
        if (!reader.Error().empty()) {
            rtReportError("OBJ", "Error: %s", reader.Error().c_str());
        }
        return RT_UNKNOWN_ERROR;
    }

    rt_temp_arena temp = rtGetTemporaryArena(NULL, 0);

    auto &attrib = reader.GetAttrib();
    auto &shapes = reader.GetShapes();
    // auto &materials = reader.GetMaterials();

    mesh_data *meshes = new mesh_data[shapes.size()];
    RT_VERIFY(meshes);

    for (size_t shape_idx = 0; shape_idx < shapes.size(); ++shape_idx) {

        rt_temp_arena rewind = rtBeginTempArena(temp.arena);

        size_t num_faces           = shapes[shape_idx].mesh.num_face_vertices.size();
        vertex *unindexed_vertices = RT_ARENA_PUSH_ARRAY(temp.arena, vertex, num_faces * 3);

        size_t index_offset = 0;
        for (size_t face_idx = 0; face_idx < shapes[shape_idx].mesh.num_face_vertices.size();
             ++face_idx) {
            size_t fv = shapes[shape_idx].mesh.num_face_vertices[face_idx];
            for (size_t vert_idx = 0; vert_idx < fv; ++vert_idx) {
                tinyobj::index_t idx = shapes[shape_idx].mesh.indices[index_offset + vert_idx];
                tinyobj::real_t vx   = attrib.vertices[3 * idx.vertex_index];
                tinyobj::real_t vy   = attrib.vertices[3 * idx.vertex_index + 1];
                tinyobj::real_t vz   = attrib.vertices[3 * idx.vertex_index + 2];

                unindexed_vertices[index_offset + vert_idx].vx = vx;
                unindexed_vertices[index_offset + vert_idx].vy = vy;
                unindexed_vertices[index_offset + vert_idx].vz = vz;
            }

            index_offset += fv;
        }

        // Generate a non-redundant index buffer
        size_t num_indices  = num_faces * 3;
        unsigned int *remap = RT_ARENA_PUSH_ARRAY(temp.arena, unsigned int, num_indices);
        size_t num_vertices = meshopt_generateVertexRemap(remap,
                                                          nullptr,
                                                          num_indices,
                                                          unindexed_vertices,
                                                          num_faces * 3,
                                                          sizeof(vertex));

        meshes[shape_idx].num_indices  = (uint32_t)num_indices;
        meshes[shape_idx].indices      = new uint32_t[num_indices];
        meshes[shape_idx].num_vertices = (uint32_t)num_vertices;
        meshes[shape_idx].vertices     = new vertex[num_vertices];

        meshopt_remapIndexBuffer(meshes[shape_idx].indices, nullptr, num_indices, remap);
        meshopt_remapVertexBuffer(meshes[shape_idx].vertices,
                                  unindexed_vertices,
                                  num_faces * 3,
                                  sizeof(vertex),
                                  remap);

        rtEndTempArena(rewind);
    }

    m_num_meshes = (uint32_t)shapes.size();
    m_meshes     = meshes;

    return RT_SUCCESS;
}

rt_result meshlet_generator::RunFlat(uint32_t mesh_idx) {
    const float cone_weight = 0.5f;

    rt_temp_arena temp = rtGetTemporaryArena(NULL, 0);

    size_t max_meshlets       = meshopt_buildMeshletsBound(m_meshes[mesh_idx].num_indices,
                                                     MESHLET_VERTICES,
                                                     MESHLET_TRIANGLES);
    meshopt_Meshlet *meshlets = RT_ARENA_PUSH_ARRAY(temp.arena, meshopt_Meshlet, max_meshlets);
    unsigned int *meshlet_vertices =
        RT_ARENA_PUSH_ARRAY(temp.arena, unsigned int, max_meshlets *MESHLET_VERTICES);
    unsigned char *meshlet_triangles =
        RT_ARENA_PUSH_ARRAY(temp.arena, unsigned char, max_meshlets *MESHLET_INDICES);

    size_t meshlet_count = meshopt_buildMeshlets(meshlets,
                                                 meshlet_vertices,
                                                 meshlet_triangles,
                                                 m_meshes[mesh_idx].indices,
                                                 m_meshes[mesh_idx].num_indices,
                                                 &m_meshes[mesh_idx].vertices[0].vx,
                                                 m_meshes[mesh_idx].num_vertices,
                                                 sizeof(vertex),
                                                 MESHLET_VERTICES,
                                                 MESHLET_TRIANGLES,
                                                 cone_weight);

    m_meshlets     = new meshlet[meshlet_count];
    m_num_meshlets = (uint32_t)meshlet_count;

    for (size_t i = 0; i < meshlet_count; ++i) {
        meshopt_optimizeMeshlet(&meshlet_vertices[meshlets[i].vertex_offset],
                                &meshlet_triangles[meshlets[i].triangle_offset],
                                meshlets[i].triangle_count,
                                meshlets[i].vertex_count);

        meshopt_Bounds bounds =
            meshopt_computeMeshletBounds(&meshlet_vertices[meshlets[i].vertex_offset],
                                         &meshlet_triangles[meshlets[i].triangle_offset],
                                         meshlets[i].triangle_count,
                                         &m_meshes[mesh_idx].vertices[0].vx,
                                         m_meshes[mesh_idx].num_vertices,
                                         sizeof(vertex));

        m_meshlets[i].vertices     = new vertex[MESHLET_VERTICES];
        m_meshlets[i].indices      = new uint8_t[MESHLET_INDICES];
        m_meshlets[i].num_vertices = meshlets[i].vertex_count;
        m_meshlets[i].num_indices  = meshlets[i].triangle_count * 3;
        memcpy(m_meshlets[i].center, bounds.center, sizeof(bounds.center));
        m_meshlets[i].radius = bounds.radius;
        memcpy(m_meshlets[i].cone_axis, bounds.cone_axis, sizeof(bounds.cone_axis));
        m_meshlets[i].cone_cutoff = bounds.cone_cutoff;
        memcpy(m_meshlets[i].cone_apex, bounds.cone_apex, sizeof(bounds.cone_apex));

        for (unsigned int vert_idx = 0; vert_idx < meshlets[i].vertex_count; ++vert_idx) {
            unsigned int vert = meshlet_vertices[meshlets[i].vertex_offset + vert_idx];
            m_meshlets[i].vertices[vert_idx] = m_meshes[mesh_idx].vertices[vert];
        }

        memcpy(m_meshlets[i].indices,
               meshlet_triangles + meshlets[i].triangle_offset,
               meshlets[i].triangle_count * 3);
        m_meshlets[i].num_children = 0u;
    }

    return RT_SUCCESS;
}

rt_result meshlet_generator::RunHierarchical(uint32_t mesh_idx, uint32_t *out_root) {
    const float cone_weight = 0.25f;

    rt_temp_arena temp = rtGetTemporaryArena(NULL, 0);

    size_t max_meshlets       = meshopt_buildMeshletsBound(m_meshes[mesh_idx].num_indices,
                                                     MESHLET_VERTICES,
                                                     MESHLET_TRIANGLES);
    meshopt_Meshlet *meshlets = RT_ARENA_PUSH_ARRAY(temp.arena, meshopt_Meshlet, max_meshlets);
    unsigned int *meshlet_vertices =
        RT_ARENA_PUSH_ARRAY(temp.arena, unsigned int, max_meshlets *MESHLET_VERTICES);
    unsigned char *meshlet_triangles =
        RT_ARENA_PUSH_ARRAY(temp.arena, unsigned char, max_meshlets *MESHLET_INDICES);
    meshopt_Bounds *meshlet_bounds = RT_ARENA_PUSH_ARRAY(temp.arena, meshopt_Bounds, max_meshlets);

    uint32_t meshlet_count = (uint32_t)meshopt_buildMeshlets(meshlets,
                                                             meshlet_vertices,
                                                             meshlet_triangles,
                                                             m_meshes[mesh_idx].indices,
                                                             m_meshes[mesh_idx].num_indices,
                                                             &m_meshes[mesh_idx].vertices[0].vx,
                                                             m_meshes[mesh_idx].num_vertices,
                                                             sizeof(vertex),
                                                             MESHLET_VERTICES,
                                                             MESHLET_TRIANGLES,
                                                             cone_weight);

    for (size_t i = 0; i < meshlet_count; ++i) {
        meshopt_optimizeMeshlet(&meshlet_vertices[meshlets[i].vertex_offset],
                                &meshlet_triangles[meshlets[i].triangle_offset],
                                meshlets[i].triangle_count,
                                meshlets[i].vertex_count);

        meshlet_bounds[i] =
            meshopt_computeMeshletBounds(&meshlet_vertices[meshlets[i].vertex_offset],
                                         &meshlet_triangles[meshlets[i].triangle_offset],
                                         meshlets[i].triangle_count,
                                         &m_meshes[mesh_idx].vertices[0].vx,
                                         m_meshes[mesh_idx].num_vertices,
                                         sizeof(vertex));
    }

    // We now have a flat list of meshlets -> the highest lod ones
    // We now combine (up to 8) meshlets into one to generate the next hierarchy level
    // Repeat until we only have 1 meshlet left
    std::vector<uint32_t> unprocessed;
    unprocessed.reserve(meshlet_count);
    for (uint32_t i = 0; i < meshlet_count; ++i) {
        unprocessed.push_back(i);
    }
    std::vector<uint32_t> next_level;
    next_level.reserve((unprocessed.size() + 7) / 8);
    while (!unprocessed.empty()) {
        uint32_t first = unprocessed.back();
        unprocessed.pop_back();

        // Find the 7 closest center points
        uint32_t closest[7];
        float distances[7];
        unsigned int closest_count = 0u;

        float first_center[3];
        memcpy(first_center, meshlet_bounds[first].center, sizeof(float) * 3);

        for (uint32_t i = 0; i < unprocessed.size(); ++i) {
            float center[3];
            memcpy(center, meshlet_bounds[unprocessed[i]].center, sizeof(float) * 3);

            float dist = sqrtf((center[0] - first_center[0]) * (center[0] - first_center[0]) +
                               (center[1] - first_center[1]) * (center[1] - first_center[1]) +
                               (center[2] - first_center[2]) * (center[2] - first_center[2]));

            if (closest_count == RT_ARRAY_COUNT(closest)) {
                // Check if we are closer than one of the other candidates
                for (unsigned int j = 0; j < closest_count; ++j) {
                    uint32_t highest_idx = UINT_MAX;
                    float highest_dist = dist;
                    if (dist < distances[j]) {
                        if (distances[j] > highest_dist) {
                            highest_dist = distances[j];
                            highest_idx  = j;
                        }
                    }
                    if (highest_idx < RT_ARRAY_COUNT(closest)) {
                        const uint32_t replaced = highest_idx;
                        distances[j]            = dist;
                        closest[j]              = i;
                        unprocessed.push_back(replaced);
                    }
                }
            } else {
                closest[closest_count]   = i;
                distances[closest_count] = dist;
                closest_count++;
                unprocessed.erase(unprocessed.begin() + i);
            }
        }

        // Combine into a new meshlet
        // vertex *vertices = new vertex[MESHLET_VERTICES * 8];
    }
    return RT_SUCCESS;
}