#include "meshlet_renderer.hpp"
#include <imgui.h>

#include <stdio.h>

#define GLM_ENABLE_EXPERIMENTAL
#include <glm/glm.hpp>
#include <glm/gtx/transform.hpp>
#include <glm/gtc/quaternion.hpp>

shader shader::CompileSource(const char *vert_src, const char *frag_src) {
    shader sh = {0};

    GLuint vert, frag;
    GLint status;
    vert = glCreateShader(GL_VERTEX_SHADER);
    glShaderSource(vert, 1, &vert_src, NULL);
    glCompileShader(vert);
    glGetShaderiv(vert, GL_COMPILE_STATUS, &status);
    if (status != GL_TRUE) {
        char log[512];
        glGetShaderInfoLog(vert, 512, NULL, log);
        rtLog("MESHLETS", "VERT: %s", log);
        return sh;
    }

    frag = glCreateShader(GL_FRAGMENT_SHADER);
    glShaderSource(frag, 1, &frag_src, NULL);
    glCompileShader(frag);
    glGetShaderiv(frag, GL_COMPILE_STATUS, &status);
    if (status != GL_TRUE) {
        char log[512];
        glGetShaderInfoLog(frag, 512, NULL, log);
        rtLog("MESHLETS", "FRAG: %s", log);
        return sh;
    }

    sh.m_prog = glCreateProgram();
    glAttachShader(sh.m_prog, vert);
    glAttachShader(sh.m_prog, frag);
    glLinkProgram(sh.m_prog);
    glGetProgramiv(sh.m_prog, GL_LINK_STATUS, &status);
    if (status != GL_TRUE) {
        char log[512];
        glGetProgramInfoLog(sh.m_prog, 512, NULL, log);
        rtLog("MESHLETS", "PROG: %s", log);
        sh.m_prog = 0;
        return sh;
    }
    glDeleteShader(vert);
    glDeleteShader(frag);

    return sh;
}

shader shader::CompileSource(const char *compute_src) {
    shader sh = {0};

    GLuint comp;
    GLint status;
    comp = glCreateShader(GL_COMPUTE_SHADER);
    glShaderSource(comp, 1, &compute_src, NULL);
    glCompileShader(comp);
    glGetShaderiv(comp, GL_COMPILE_STATUS, &status);
    if (status != GL_TRUE) {
        char log[512];
        glGetShaderInfoLog(comp, 512, NULL, log);
        rtLog("MESHLETS", "COMP: %s", log);
        return sh;
    }

    sh.m_prog = glCreateProgram();
    glAttachShader(sh.m_prog, comp);
    glLinkProgram(sh.m_prog);
    glGetProgramiv(sh.m_prog, GL_LINK_STATUS, &status);
    if (status != GL_TRUE) {
        char log[512];
        glGetProgramInfoLog(sh.m_prog, 512, NULL, log);
        rtLog("MESHLETS", "PROG: %s", log);
        sh.m_prog = 0;
        return sh;
    }
    glDeleteShader(comp);
    return sh;
}

shader shader::CompileFile(const char *compute_path) {
    shader sh = {0};
    FILE *f   = fopen(compute_path, "rb");
    if (!f)
        return sh;
    fseek(f, 0, SEEK_END);
    long fsz = ftell(f);
    fseek(f, 0, SEEK_SET);
    char *buf = new char[fsz + 1];
    fread(buf, 1, fsz, f);
    buf[fsz] = 0;
    fclose(f);

    sh = CompileSource(buf);
    delete buf;
    return sh;
}

struct gpu_mesh_data {
    glm::mat4 model;
};

struct gpu_meshlet_data {
    float bounds[4];    // xyz, radius
    float cone_apex[4]; // xyz, unused
    float cone_axis_cutoff[4];
    uint32_t mesh_index;
    uint32_t first_index;
    uint32_t base_vertex;
    uint32_t index_count;
};

struct draw_elements_indirect_command {
    uint32_t count;
    uint32_t instanceCount;
    uint32_t firstIndex;
    int32_t baseVertex;
    uint32_t baseInstance;
};

struct cull_output {
    uint32_t draw_count;
};

static const char *_single_vert = "#version 460\n"
                                  "layout (location = 0) in vec3 vertpos;\n"
                                  "uniform mat4 mvp;\n"
                                  "void main() {\n"
                                  "  gl_Position = mvp * vec4(vertpos, 1.0);\n"
                                  "}\n";

static const char *_single_frag = "#version 460\n"
                                  "out vec3 fragcol;\n"
                                  "uniform vec3 meshletcol;\n"
                                  "void main(){\n"
                                  "  fragcol = meshletcol;\n"
                                  "}\n";

static const char *_flat_cull = "src/experimental/meshlets/flat_cull.glsl";

static const char *_meshlet_vert = "#version 460\n"
                                   "layout (location = 0) in vec3 vertpos;\n"
                                   "uniform mat4 mvp;\n"
                                   "out vec3 meshletcol;\n"
                                   "vec3 colors[5] = vec3[5](\n"
                                   "  vec3(129.f / 255.f, 132.f / 255.f, 121.f / 255.f),\n"
                                   "  vec3(181.f / 255.f, 203.f / 255.f, 183.f / 255.f),\n"
                                   "  vec3(210.f / 255.f, 228.f / 255.f, 196.f / 255.f),\n"
                                   "  vec3(228.f / 255.f, 233.f / 255.f, 178.f / 255.f),\n"
                                   "  vec3(231.f / 255.f, 224.f / 255.f, 139 / 255.f));\n"
                                   "void main() {\n"
                                   "  gl_Position = mvp * vec4(vertpos, 1.0);\n"
                                   "  meshletcol = colors[gl_DrawID % 5];\n"
                                   "}\n";

static const char *_meshlet_frag = "#version 460\n"
                                   "out vec3 fragcol;\n"
                                   "in vec3 meshletcol;\n"
                                   "void main(){\n"
                                   "  fragcol = meshletcol;\n"
                                   "}\n";

rt_result meshlet_renderer::Initialize() {
    m_settings.eye[0] = 0.f;
    m_settings.eye[1] = 0.5f;
    m_settings.eye[2] = 0.5f;
    m_settings.fov    = glm::radians(45.f);
    memset(m_settings.target, 0, sizeof(m_settings.target));

    m_single_meshlet_shader = shader::CompileSource(_single_vert, _single_frag);
    m_flat_cull_shader      = shader::CompileFile(_flat_cull);
    m_meshlet_shader        = shader::CompileSource(_meshlet_vert, _meshlet_frag);
    return RT_SUCCESS;
}

void meshlet_renderer::RenderFlat(const meshlet *meshlets, unsigned int count) {
    if (m_settings.separate_rendering) {
        SeparateRendering(meshlets, count);
    } else {
        DrawIndirectFlat(meshlets, count);
    }
}

void meshlet_renderer::DrawIndirectFlat(const meshlet *meshlets, unsigned int count) {
    // Store all meshlets in one vertex- and index buffer
    // Store meshlet metadata in one ssbo
    // Do compute "culling" (generate drawindirect) into 1 ssbo
    // DrawIndirect

    GLuint vbo, ebo, meshlet_ssbo, draw_ssbo, cull_ssbo, mesh_ssbo;
    GLuint buffers[6];
    glCreateBuffers(RT_ARRAY_COUNT(buffers), buffers);
    vbo          = buffers[0];
    ebo          = buffers[1];
    meshlet_ssbo = buffers[2];
    draw_ssbo    = buffers[3];
    cull_ssbo    = buffers[4];
    mesh_ssbo    = buffers[5];

    // Store vertices and indices at idx * MESHLET_VERTICES/MESHLET_INDICES
    glNamedBufferStorage(vbo,
                         count * sizeof(vertex) * MESHLET_VERTICES,
                         nullptr,
                         GL_DYNAMIC_STORAGE_BIT);
    glNamedBufferStorage(ebo, count * MESHLET_INDICES, nullptr, GL_DYNAMIC_STORAGE_BIT);
    for (unsigned int i = 0; i < count; ++i) {
        glNamedBufferSubData(vbo,
                             i * sizeof(vertex) * MESHLET_VERTICES,
                             meshlets[i].num_vertices * sizeof(vertex),
                             meshlets[i].vertices);
        glNamedBufferSubData(ebo,
                             i * MESHLET_INDICES,
                             meshlets[i].num_indices,
                             meshlets[i].indices);
    }

    // Store meshlet information
    glNamedBufferStorage(meshlet_ssbo,
                         count * sizeof(gpu_meshlet_data),
                         nullptr,
                         GL_DYNAMIC_STORAGE_BIT);
    for (unsigned int i = 0; i < count; ++i) {
        gpu_meshlet_data meshlet;
        // Only have one right now
        meshlet.mesh_index  = 0;
        meshlet.first_index = i * MESHLET_INDICES;
        meshlet.base_vertex = i * MESHLET_VERTICES;
        meshlet.index_count = meshlets[i].num_indices;
        memcpy(meshlet.bounds, meshlets[i].center, 3 * sizeof(float));
        meshlet.bounds[3] = meshlets[i].radius;
        memcpy(meshlet.cone_apex, meshlets[i].cone_apex, 3 * sizeof(float));
        meshlet.cone_apex[3] = 0.f;
        memcpy(meshlet.cone_axis_cutoff, meshlets[i].cone_axis, sizeof(meshlets[i].cone_axis));
        glNamedBufferSubData(meshlet_ssbo,
                             i * sizeof(gpu_meshlet_data),
                             sizeof(gpu_meshlet_data),
                             &meshlet);
    }

    // Reserve space for the draw commands
    glNamedBufferStorage(draw_ssbo, count * sizeof(draw_elements_indirect_command), nullptr, 0);

    // Prepare culling output
    cull_output cull_output = {0};
    glNamedBufferStorage(cull_ssbo, sizeof(cull_output), &cull_output, GL_MAP_READ_BIT);

    // Prepare mesh data
    gpu_mesh_data mesh_data;
    mesh_data.model = glm ::rotate(m_settings.rotation, glm::vec3(0, 1, 0)) *
                      glm::scale(glm::vec3(m_settings.scale));
    glNamedBufferStorage(mesh_ssbo, sizeof(mesh_data), &mesh_data, GL_DYNAMIC_STORAGE_BIT);

    glMemoryBarrier(GL_BUFFER_UPDATE_BARRIER_BIT);

    // Do culling. TODO: Get number of draws back
    m_flat_cull_shader.Use();
    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, meshlet_ssbo);
    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, draw_ssbo);
    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 3, cull_ssbo);
    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 4, mesh_ssbo);
    glUniform3fv(0, 1, m_settings.eye);
    glDispatchCompute(count, 1, 1);

    // Create the vao
    GLuint vao;
    glGenVertexArrays(1, &vao);
    glBindVertexArray(vao);
    glBindBuffer(GL_ARRAY_BUFFER, vbo);
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ebo);
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, sizeof(vertex), nullptr);
    glEnableVertexAttribArray(0);

    // DrawIndirect
    GLuint mvploc = glGetUniformLocation(m_meshlet_shader.m_prog, "mvp");

    glm::mat4 model = glm::rotate(m_settings.rotation, glm::vec3(0, 1, 0)) *
                      glm::scale(glm::vec3(m_settings.scale));
    glm::mat4 view =
        glm::lookAt(glm::vec3(m_settings.eye[0], m_settings.eye[1], m_settings.eye[2]),
                    glm::vec3(m_settings.target[0], m_settings.target[1], m_settings.target[2]),
                    glm::vec3(0, 1, 0));
    glm::mat4 proj = glm::perspective(m_settings.fov, m_aspect, 0.01f, 100.f);
    glm::mat4 mvp  = proj * view * model;

    glDisable(GL_CULL_FACE);
    glCullFace(GL_BACK);
    glEnable(GL_DEPTH_TEST);

    m_meshlet_shader.Use();
    glUniformMatrix4fv(mvploc, 1, GL_FALSE, &mvp[0][0]);

    glBindBuffer(GL_DRAW_INDIRECT_BUFFER, draw_ssbo);
    glMemoryBarrier(GL_COMMAND_BARRIER_BIT | GL_SHADER_STORAGE_BARRIER_BIT);
    {
        void *out = glMapNamedBuffer(cull_ssbo, GL_READ_ONLY);
        memcpy(&cull_output, out, sizeof(cull_output));
        glUnmapNamedBuffer(cull_ssbo);
    }

    glMultiDrawElementsIndirect(GL_TRIANGLES,
                                GL_UNSIGNED_BYTE,
                                nullptr,
                                cull_output.draw_count,
                                sizeof(draw_elements_indirect_command));

    glDeleteBuffers(RT_ARRAY_COUNT(buffers), buffers);
}

void meshlet_renderer::SeparateRendering(const meshlet *meshlets, unsigned int count) {
    GLuint vao, vbo, ebo;
    glGenBuffers(1, &vbo);
    glGenBuffers(1, &ebo);
    glGenVertexArrays(1, &vao);

    glBindVertexArray(vao);
    glBindBuffer(GL_ARRAY_BUFFER, vbo);
    glBufferData(GL_ARRAY_BUFFER, sizeof(vertex) * MESHLET_VERTICES, nullptr, GL_STREAM_DRAW);
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ebo);
    glBufferData(GL_ELEMENT_ARRAY_BUFFER, MESHLET_INDICES, nullptr, GL_STREAM_DRAW);

    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, sizeof(vertex), nullptr);
    glEnableVertexAttribArray(0);

    m_single_meshlet_shader.Use();

    float colors[5][3] = {
        {129.f / 255.f, 132.f / 255.f, 121.f / 255.f},
        {181.f / 255.f, 203.f / 255.f, 183.f / 255.f},
        {210.f / 255.f, 228.f / 255.f, 196.f / 255.f},
        {228.f / 255.f, 233.f / 255.f, 178.f / 255.f},
        {231.f / 255.f, 224.f / 255.f,   139 / 255.f},
    };

    GLuint colorloc = glGetUniformLocation(m_single_meshlet_shader.m_prog, "meshletcol");
    GLuint mvploc   = glGetUniformLocation(m_single_meshlet_shader.m_prog, "mvp");

    glm::mat4 model = glm::rotate(m_settings.rotation, glm::vec3(0, 1, 0)) *
                      glm::scale(glm::vec3(m_settings.scale));
    glm::mat4 view =
        glm::lookAt(glm::vec3(m_settings.eye[0], m_settings.eye[1], m_settings.eye[2]),
                    glm::vec3(m_settings.target[0], m_settings.target[1], m_settings.target[2]),
                    glm::vec3(0, 1, 0));
    glm::mat4 proj = glm::perspective(m_settings.fov, m_aspect, 0.01f, 100.f);
    glm::mat4 mvp  = proj * view * model;

    glUniformMatrix4fv(mvploc, 1, GL_FALSE, &mvp[0][0]);
    glDisable(GL_CULL_FACE);
    glCullFace(GL_BACK);
    glEnable(GL_DEPTH_TEST);

    for (unsigned int i = 0; i < count; ++i) {
        glBufferSubData(GL_ARRAY_BUFFER,
                        0,
                        sizeof(vertex) * meshlets[i].num_vertices,
                        meshlets[i].vertices);
        glBufferSubData(GL_ELEMENT_ARRAY_BUFFER,
                        0,
                        sizeof(uint8_t) * meshlets[i].num_indices,
                        meshlets[i].indices);
        glUniform3fv(colorloc, 1, colors[i % 5]);

        glDrawElements(GL_TRIANGLES, meshlets[i].num_indices, GL_UNSIGNED_BYTE, nullptr);
    }

    glDeleteVertexArrays(1, &vao);
    glDeleteBuffers(1, &vbo);
    glDeleteBuffers(1, &ebo);
}

void meshlet_renderer::SettingMenu() {
    settings *sett = &m_settings;
    if (ImGui::Begin("Settings")) {
        ImGui::Checkbox("Separate Rendering", &sett->separate_rendering);
        ImGui::SeparatorText("Model");
        ImGui::SliderFloat("Scale", &sett->scale, 1.f, 100.f);
        ImGui::SliderAngle("Rotation", &sett->rotation, 0.f, 360.f);
        ImGui::SeparatorText("Camera");
        ImGui::SliderAngle("Fov", &sett->fov, 30.f, 180.f);
        ImGui::InputFloat3("Eye", sett->eye);
        ImGui::InputFloat3("Target", sett->target);
        ImGui::End();
    }
}