#include "renderer/gl_renderer.hpp"

#include "components/camera.hpp"
#include "core/io.hpp"
#include "graphics/opengl/buffer.hpp"
#include "graphics/opengl/render_context.hpp"
#include "graphics/opengl/shader.hpp"
#include "graphics/opengl/vertex_array.hpp"
#include "maths/aspect_ratio.hpp"
#include "resources/primitive.hpp"
#include "utils/assert.hpp"
#include "utils/colours.hpp"

// These headers are generated using bin2hpp during build
#include "clustering_cs.hpp"
#include "default_fs.hpp"
#include "default_vs.hpp"
#include "light_culling_cs.hpp"

#include <glad/gl.h>

#include <algorithm>
#include <cstdint>
#include <glm/ext/matrix_clip_space.hpp>
#include <glm/ext/matrix_transform.hpp>
#include <glm/geometric.hpp>
#include <glm/gtc/type_ptr.hpp>
#include <glm/mat3x3.hpp>
#include <glm/mat4x4.hpp>
#include <glm/matrix.hpp>
#include <glm/trigonometric.hpp>
#include <glm/vec2.hpp>
#include <glm/vec3.hpp>

#include <cstddef>
#include <numeric>
#include <string_view>

using namespace kuiper::gl;

glm::mat4 camera_to_view_matrix(const glm::vec3& pos, const kuiper::component::camera_props& props) noexcept;

renderer renderer::init() {
    gl_state state {};

    // Configuration

    state.max_texture_units = render_context::get<GLint>(GL_MAX_TEXTURE_IMAGE_UNITS);
    state.texture_unit_slots.resize(state.max_texture_units);
    std::iota(state.texture_unit_slots.begin(), state.texture_unit_slots.end(), 0);

    // Buffers

    state.vertex_buffer = buffer::make(GL_ARRAY_BUFFER);
    state.index_buffer  = buffer::make(GL_ELEMENT_ARRAY_BUFFER);

    state.vao = vertex_array::make();
    state.vao->bind();

    constexpr std::uint32_t pos_idx        = 0;
    constexpr std::uint32_t normal_idx     = 1;
    constexpr std::uint32_t tangent_idx    = 2;
    constexpr std::uint32_t uv_idx         = 3;
    constexpr std::uint32_t draw_id_idx    = 4;
    constexpr std::uint32_t model_mat_idx  = 5; // mat4 occupies slot 5, 6, 7 & 8
    constexpr std::uint32_t normal_mat_idx = 9;

    state.vao->enable_attrib(pos_idx);
    state.vao->enable_attrib(normal_idx);
    state.vao->enable_attrib(tangent_idx);
    state.vao->enable_attrib(uv_idx);
    state.vao->enable_attrib(draw_id_idx);
    state.vao->enable_attrib(model_mat_idx + 0);
    state.vao->enable_attrib(model_mat_idx + 1);
    state.vao->enable_attrib(model_mat_idx + 2);
    state.vao->enable_attrib(model_mat_idx + 3);
    state.vao->enable_attrib(normal_mat_idx + 0);
    state.vao->enable_attrib(normal_mat_idx + 1);
    state.vao->enable_attrib(normal_mat_idx + 2);

    state.index_buffer->bind();

    state.vertex_buffer->bind();
    state.vertex_buffer->set_vertex_attrib(
        pos_idx, glm::vec3::length(), GL_FLOAT, sizeof(renderer::vertex_t), offsetof(renderer::vertex_t, position));
    state.vertex_buffer->set_vertex_attrib(normal_idx,
                                           glm::vec3::length(),
                                           GL_FLOAT,
                                           sizeof(renderer::vertex_t),
                                           offsetof(renderer::vertex_t, normal),
                                           true);
    state.vertex_buffer->set_vertex_attrib(
        tangent_idx, glm::vec3::length(), GL_FLOAT, sizeof(renderer::vertex_t), offsetof(renderer::vertex_t, tangent));
    state.vertex_buffer->set_vertex_attrib(
        uv_idx, glm::vec2::length(), GL_FLOAT, sizeof(renderer::vertex_t), offsetof(renderer::vertex_t, uv), true);

    state.matrix_buffer = buffer::make(GL_ARRAY_BUFFER);
    state.matrix_buffer->bind();
    // mat4 is essentially four vec4s
    state.matrix_buffer->set_vertex_attrib(model_mat_idx + 0,
                                           glm::vec4::length(),
                                           GL_FLOAT,
                                           sizeof(render_matrices),
                                           offsetof(render_matrices, model_matrix) + (0 * sizeof(glm::vec4)));
    state.matrix_buffer->set_vertex_attrib(model_mat_idx + 1,
                                           glm::vec4::length(),
                                           GL_FLOAT,
                                           sizeof(render_matrices),
                                           offsetof(render_matrices, model_matrix) + (1 * sizeof(glm::vec4)));
    state.matrix_buffer->set_vertex_attrib(model_mat_idx + 2,
                                           glm::vec4::length(),
                                           GL_FLOAT,
                                           sizeof(render_matrices),
                                           offsetof(render_matrices, model_matrix) + (2 * sizeof(glm::vec4)));
    state.matrix_buffer->set_vertex_attrib(model_mat_idx + 3,
                                           glm::vec4::length(),
                                           GL_FLOAT,
                                           sizeof(render_matrices),
                                           offsetof(render_matrices, model_matrix) + (3 * sizeof(glm::vec4)));
    state.matrix_buffer->set_vertex_attrib(normal_mat_idx + 0,
                                           glm::vec3::length(),
                                           GL_FLOAT,
                                           sizeof(render_matrices),
                                           offsetof(render_matrices, normal_matrix) + (0 * sizeof(glm::vec3)));
    state.matrix_buffer->set_vertex_attrib(normal_mat_idx + 1,
                                           glm::vec3::length(),
                                           GL_FLOAT,
                                           sizeof(render_matrices),
                                           offsetof(render_matrices, normal_matrix) + (1 * sizeof(glm::vec3)));
    state.matrix_buffer->set_vertex_attrib(normal_mat_idx + 2,
                                           glm::vec3::length(),
                                           GL_FLOAT,
                                           sizeof(render_matrices),
                                           offsetof(render_matrices, normal_matrix) + (2 * sizeof(glm::vec3)));

    state.indirect_buffer = buffer::make(GL_DRAW_INDIRECT_BUFFER);
    glBindBuffer(GL_ARRAY_BUFFER, state.indirect_buffer->id());
    glVertexAttribIPointer(draw_id_idx,
                           1,
                           GL_UNSIGNED_INT,
                           sizeof(draw_elements_indirect_cmd),
                           reinterpret_cast<const void*>(offsetof(draw_elements_indirect_cmd, base_instance)));

    state.vao->set_attrib_divisor(pos_idx, 0); // every fragment
    state.vao->set_attrib_divisor(normal_idx, 0);
    state.vao->set_attrib_divisor(tangent_idx, 0);
    state.vao->set_attrib_divisor(uv_idx, 0);
    state.vao->set_attrib_divisor(draw_id_idx, 1); // every instance
    state.vao->set_attrib_divisor(model_mat_idx + 0, 1);
    state.vao->set_attrib_divisor(model_mat_idx + 1, 1);
    state.vao->set_attrib_divisor(model_mat_idx + 2, 1);
    state.vao->set_attrib_divisor(model_mat_idx + 3, 1);
    state.vao->set_attrib_divisor(normal_mat_idx + 0, 1);
    state.vao->set_attrib_divisor(normal_mat_idx + 1, 1);
    state.vao->set_attrib_divisor(normal_mat_idx + 2, 1);

    glBindVertexArray(0);

    state.cluster_buffer = buffer::make(GL_SHADER_STORAGE_BUFFER);
    state.cluster_buffer->upload(GL_STATIC_COPY, NUM_CLUSTERS * sizeof(cluster), nullptr);

    state.lights_buffer = buffer::make(GL_SHADER_STORAGE_BUFFER);

    // Shaders

    auto vertex = shader::from_source(bin::default_vs, shader_type::vertex);

    // We pre-process the fragment shader to set the correct number of texture unit slots on the system
    std::string                fragment_src(bin::default_fs);
    constexpr std::string_view fs_search_and_replace {"const uint MAX_FRAG_TEX_UNITS = 32u;"};
    fragment_src.replace(fragment_src.find(fs_search_and_replace),
                         fs_search_and_replace.length(),
                         std::format("const uint MAX_FRAG_TEX_UNITS = {}u;", state.max_texture_units));
    auto fragment = shader::from_source(fragment_src, shader_type::fragment);

    state.pbr_shader = program::from_shaders({vertex, fragment});
    vertex->destroy();
    fragment->destroy();

    auto clustering_shader  = shader::from_source(bin::clustering_cs, shader_type::compute);
    state.clustering_shader = program::from_shaders({clustering_shader});
    clustering_shader->destroy();

    auto light_culling_shader  = shader::from_source(bin::light_culling_cs, shader_type::compute);
    state.light_culling_shader = program::from_shaders({light_culling_shader});
    light_culling_shader->destroy();

    return {std::move(state)};
}

void renderer::render(const framebuffer::s_ptr&      framebuffer,
                      const glm::vec3&               cam_pos,
                      const component::camera_props& cam_props) noexcept {
    update_and_upload();

    // Framebuffer

    const auto& framebuffer_size = framebuffer->dimensions();
    float       fb_x             = static_cast<float>(framebuffer_size.x);
    float       fb_y             = static_cast<float>(framebuffer_size.y);

    framebuffer->bind();
    render_context::set_viewport_size(framebuffer_size.x, framebuffer_size.y);

    glClearColor(colours::black.r, colours::black.g, colours::black.b, colours::black.a);
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

    // View & projection matrices

    const auto view_matrix = camera_to_view_matrix(cam_pos, cam_props);
    const auto projection_matrix =
        glm::perspective(cam_props.fov, maths::aspect_ratio<float>(fb_x, fb_y), NEAR_Z_DEPTH, FAR_Z_DEPTH);

    // Compute shaders: clustering & light culling

    constexpr glm::uvec3 grid_size_vec {NUM_CLUSTERS_X, NUM_CLUSTERS_Y, NUM_CLUSTERS_Z};

    m_state.clustering_shader->use();
    m_state.clustering_shader->set_float("z_near", NEAR_Z_DEPTH);
    m_state.clustering_shader->set_float("z_far", FAR_Z_DEPTH);
    m_state.clustering_shader->set_mat4("inverse_proj_mat", glm::inverse(projection_matrix));
    m_state.clustering_shader->set_uvec3("grid_size", grid_size_vec);
    m_state.clustering_shader->set_uvec2("screen_dimensions", framebuffer_size);

    m_state.cluster_buffer->bind_to_location(1);
    glDispatchCompute(NUM_CLUSTERS_X, NUM_CLUSTERS_Y, NUM_CLUSTERS_Z);
    glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

    m_state.light_culling_shader->use();
    m_state.light_culling_shader->set_mat4("view_mat", view_matrix);

    m_state.cluster_buffer->bind_to_location(1);
    m_state.lights_buffer->bind_to_location(2);

    glDispatchCompute(NUM_CLUSTERS / LOCAL_SIZE, 1, 1);
    glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

    // Geometry & shading

    m_state.vao->bind();

    m_state.pbr_shader->use();
    m_state.pbr_shader->set_mat4("view_mat", view_matrix);
    m_state.pbr_shader->set_mat4("view_proj_mat", projection_matrix * view_matrix);

    m_state.pbr_shader->set_float("z_near", NEAR_Z_DEPTH);
    m_state.pbr_shader->set_float("z_far", FAR_Z_DEPTH);
    m_state.pbr_shader->set_uvec3("grid_size", grid_size_vec);
    m_state.pbr_shader->set_uvec2("screen_dimensions", framebuffer_size);
    m_state.pbr_shader->set_vec3("view_pos", cam_pos);

    m_state.pbr_shader->set_int_array(
        "packed_tex", m_state.texture_unit_slots.data(), m_state.texture_unit_slots.size());

    for (std::uint32_t draw_count = 0; draw_count < m_draw_cmds.size(); draw_count += m_state.max_texture_units) {
        for (std::uint32_t i = 0; i < m_state.max_texture_units && i + draw_count < m_materials.size(); ++i) {
            m_materials[i + draw_count].textures->bind(i);
        }

        glMultiDrawElementsIndirect(
            GL_TRIANGLES,
            GL_UNSIGNED_INT,
            reinterpret_cast<const void*>(draw_count * sizeof(draw_elements_indirect_cmd)),
            static_cast<GLsizei>(
                std::min(static_cast<std::uint32_t>(m_draw_cmds.size() - draw_count), m_state.max_texture_units)),
            sizeof(draw_elements_indirect_cmd));
    }

    return;
}

void renderer::flush() noexcept {
    m_draw_cmds.clear();
    m_matrices.clear();
    m_vertices.clear();
    m_indices.clear();
    m_materials.clear();
    m_lights.clear();

    return;
}

renderer_metrics renderer::metrics() const noexcept {
    renderer_metrics metrics {};

    metrics.n_vertices  = m_vertices.size();
    metrics.n_indices   = m_indices.size();
    metrics.n_draw_cmds = m_draw_cmds.size();
    metrics.n_matrices  = m_matrices.size();
    metrics.n_materials = m_materials.size();
    metrics.n_lights    = m_lights.size();

    return metrics;
}

void renderer::update_and_upload() noexcept {
    const auto vertices_size_bytes      = m_vertices.size() * sizeof(vertex);
    const auto indices_size_bytes       = m_indices.size() * sizeof(renderer::index_t);
    const auto lights_size_bytes        = m_lights.size() * sizeof(renderer::point_light);
    const auto indirect_cmds_size_bytes = m_draw_cmds.size() * sizeof(draw_elements_indirect_cmd);
    const auto matrices_size_bytes      = m_matrices.size() * sizeof(render_matrices);

    // NOTE: assuming the use of mutable buffers here

    if (m_vertices.size() > m_state.vertex_buffer->size() / sizeof(vertex)) {
        m_state.vertex_buffer->upload(GL_STREAM_DRAW, vertices_size_bytes, m_vertices.data());
    } else {
        m_state.vertex_buffer->update(0, vertices_size_bytes, m_vertices.data());
    }

    if (m_indices.size() > m_state.index_buffer->size() / sizeof(renderer::index_t)) {
        m_state.index_buffer->upload(GL_STREAM_DRAW, indices_size_bytes, m_indices.data());
    } else {
        m_state.index_buffer->update(0, indices_size_bytes, m_indices.data());
    }

    if (m_lights.size() > m_state.lights_buffer->size() / sizeof(renderer::point_light)) {
        m_state.lights_buffer->upload(GL_STREAM_DRAW, lights_size_bytes, m_lights.data());
    } else {
        m_state.lights_buffer->update(0, lights_size_bytes, m_lights.data());
    }

    if (m_draw_cmds.size() > m_state.indirect_buffer->size() / sizeof(draw_elements_indirect_cmd)) {
        m_state.indirect_buffer->upload(GL_STREAM_DRAW, indirect_cmds_size_bytes, m_draw_cmds.data());
    } else {
        m_state.indirect_buffer->update(0, indirect_cmds_size_bytes, m_draw_cmds.data());
    }

    if (m_matrices.size() > m_state.matrix_buffer->size() / sizeof(render_matrices)) {
        m_state.matrix_buffer->upload(GL_STREAM_DRAW, matrices_size_bytes, m_matrices.data());
    } else {
        m_state.matrix_buffer->update(0, matrices_size_bytes, m_matrices.data());
    }

    return;
}

void renderer::draw_model(const resource::model& model, const maths::transform_euler& transform) {
    const auto transform_matrix = transform.to_mat4();

    for (const auto& mesh : model.meshes()) {
        draw_mesh(mesh, model.materials(), transform_matrix);
    }

    return;
}

void renderer::draw_mesh(const resource::mesh&                  mesh,
                         const std::vector<resource::material>& materials,
                         const glm::mat4&                       transform) {
    // From the glTF 2.0 spec:
    // "When the scale is zero on all three axes [...] implementations are free
    // to optimize away rendering of the node's mesh and all of the node's
    // children's meshes"
    // TODO: optimise for this ^

    // The `transform` parameter is the parent mesh node's local transformation matrix

    const auto rel_mat = transform * mesh.transform();

    for (const auto& prim : mesh.primitives()) {
        draw_primitive(prim, materials[prim.material_idx], rel_mat);
    }

    for (const auto& submesh : mesh.submeshes()) {
        draw_mesh(submesh, materials, rel_mat);
    }

    return;
}

void renderer::draw_primitive(const resource::primitive&    primitive,
                              const resource::material&     material,
                              const maths::transform_euler& transform) {
    return draw_primitive(primitive, material, transform.to_mat4());
}

void renderer::draw_primitive(const resource::primitive& primitive,
                              const resource::material&  material,
                              const glm::mat4&           transform) {
    const std::uint32_t n_vertices = primitive.vertices.size();
    const std::uint32_t n_indices  = primitive.indices.size();

    const std::uint32_t vtx_start =
        static_cast<std::uint32_t>(std::distance(std::begin(m_vertices), std::end(m_vertices)));
    const std::uint32_t idx_start =
        static_cast<std::uint32_t>(std::distance(std::begin(m_indices), std::end(m_indices)));

    KUIPER_ASSERT(vtx_start <= (std::numeric_limits<std::uint32_t>::max() - n_vertices));
    KUIPER_ASSERT(idx_start <= (std::numeric_limits<std::uint32_t>::max() - n_indices));

    draw_elements_indirect_cmd cmd {};
    cmd.count          = n_indices;
    cmd.instance_count = 1;
    cmd.first_index    = idx_start;
    cmd.base_vertex    = vtx_start;
    cmd.base_instance  = m_draw_cmds.size();

    m_vertices.insert(std::end(m_vertices), std::begin(primitive.vertices), std::end(primitive.vertices));
    m_indices.insert(std::end(m_indices), std::begin(primitive.indices), std::end(primitive.indices));

    m_matrices.push_back({transform, glm::transpose(glm::inverse(glm::mat3(transform)))});
    m_draw_cmds.push_back(cmd);
    m_materials.push_back(material);

    return;
}

void renderer::add_light(const glm::vec3& pos, const component::point_light& light) noexcept {
    m_lights.emplace_back(glm::make_vec4(pos), glm::make_vec4(light.colour), light.intensity, 50.0f);
    return;
}

glm::mat4 camera_to_view_matrix(const glm::vec3& pos, const kuiper::component::camera_props& props) noexcept {
    static constexpr glm::vec3 up_unit_vec {0.0f, 1.0f, 0.0f};

    const glm::vec3 direction = {glm::cos(props.euler.y) * glm::cos(props.euler.x),
                                 glm::sin(props.euler.x),
                                 glm::sin(props.euler.y) * glm::cos(props.euler.x)};
    const glm::vec3 cam_front = glm::normalize(direction);

    return glm::lookAt(pos, pos + cam_front, up_unit_vec);
}
