#include <regex>
#include <iostream>
#include "render.h"
#include "glm/vec3.hpp"
#include "glm/vec4.hpp"
#include "glm/gtc/type_ptr.hpp"
#include "gcode_parser.h"
#include <fstream>

static const char *vs_code = R"(
    #version 330

    layout (location = 0) in vec3 position;
    layout (location = 1) in vec4 in_color;

    out vec4 interp_color;

    out vec3 vertPos;
    flat out vec3 startPos;
    uniform mat4 mvp;


    void main() {
        vec4 pos = mvp * vec4(position, 1.0);
        gl_Position = pos;

        vertPos      = pos.xyz / pos.w;
        startPos     = vertPos;
        interp_color = in_color;
    }
)";


static const char *ps_code = R"(
    #version 330

    flat in vec3 startPos;
    in vec3 vertPos;

    in vec4 interp_color;
    out vec4 color;

    uniform vec2  u_resolution;
    float u_dashSize = 10; // make this uniform
    float u_gapSize = 10; // make this uniform

    void main() {
        vec2  dir  = (vertPos.xy - startPos.xy) * u_resolution/2.0;
        float dist = length(dir);

        if (fract(dist / (u_dashSize + u_gapSize)) > u_dashSize/(u_dashSize + u_gapSize))
            discard;

        color = interp_color;
    }
)";

static const char *vs_heightmap_code = R"(
    #version 330

    layout (location = 0) in vec3 position;
    layout (location = 1) in vec3 normal;
    layout (location = 2) in vec4 color;

    uniform mat4 mmtx;
    uniform mat4 vmtx;
    uniform mat4 pmtx;
    uniform mat3 nmtx;

    out vec3 eye_space_normal;
    out vec3 eye_space_position;
    out vec4 diffuse_color;

    void main() {
        mat4 mvp = pmtx * vmtx * mmtx;
        gl_Position = mvp * vec4(position, 1.0);

        eye_space_position = (vmtx * mmtx * vec4(position, 1)).xyz;
        eye_space_normal = normalize(nmtx * normal);

        diffuse_color = color;
    }
)";


static const char *ps_heightmap_code = R"(
    #version 330

    in vec3 eye_space_normal;
    in vec3 eye_space_position;
    in vec4 diffuse_color;


    out vec4 color;

    void main() {
        vec3 normal = normalize(eye_space_normal);
        vec3 light_dir = normalize(vec3(0, 0, 1));
        vec4 ambient_color = vec4(0.1, 0.1, 0.1, 1);

        float n_dot_l = max(dot(normal, light_dir), 0.0);
        color = n_dot_l * diffuse_color + ambient_color;
    }
)";

static void add_vertex(std::vector<float> &buffer_data, glm::vec3 v, glm::vec3 n, glm::vec4 col) {
    buffer_data.push_back(v.x);
    buffer_data.push_back(v.y);
    buffer_data.push_back(v.z);

    buffer_data.push_back(n.x);
    buffer_data.push_back(n.y);
    buffer_data.push_back(n.z);

    buffer_data.push_back(col.r);
    buffer_data.push_back(col.g);
    buffer_data.push_back(col.b);
    buffer_data.push_back(col.a);
}

static void add_vertex(std::vector<float> &buffer_data, glm::vec3 v, glm::vec4 col) {
    buffer_data.push_back(v.x);
    buffer_data.push_back(v.y);
    buffer_data.push_back(v.z);

    buffer_data.push_back(col.r);
    buffer_data.push_back(col.g);
    buffer_data.push_back(col.b);
    buffer_data.push_back(col.a);
}

static void
add_line(std::vector<float> &buffer_data, glm::vec3 from, glm::vec4 from_col, glm::vec3 to, glm::vec4 to_col) {
    add_vertex(buffer_data, from, from_col);
    add_vertex(buffer_data, to, to_col);
}

static void add_line(std::vector<float> &buffer_data, glm::vec3 from, glm::vec3 to, glm::vec4 col) {
    add_vertex(buffer_data, from, col);
    add_vertex(buffer_data, to, col);
}

static void
add_triangle(std::vector<float> &buffer_data, glm::vec3 p1, glm::vec3 p2, glm::vec3 p3, glm::vec3 n, glm::vec4 col) {
    add_vertex(buffer_data, p1, n, col);
    add_vertex(buffer_data, p2, n, col);
    add_vertex(buffer_data, p3, n, col);
}


void grbl::program_renderer::render(glm::mat4 model, glm::mat4 view, glm::mat4 projection, glm::mat3 normal_mat,
                                    glm::vec2 viewport_size) {
    if (shader == nullptr || heightmap_shader == nullptr) return;

    // draw heightmap
    heightmap_shader->bind();
    heightmap_shader->set_mat4(glm::value_ptr(model), "mmtx");
    heightmap_shader->set_mat4(glm::value_ptr(view), "vmtx");
    heightmap_shader->set_mat4(glm::value_ptr(projection), "pmtx");
    heightmap_shader->set_mat3(glm::value_ptr(normal_mat), "nmtx");

    glBindVertexArray(heightmap_vao_id);
    glDrawArrays(GL_TRIANGLES, 0, heightmap_vertices_count);

    heightmap_shader->unbind();

    auto mvp = projection * view * model;

    // draw model
    shader->bind();
    shader->set_mat4(glm::value_ptr(mvp), "mvp");
    shader->set_vec2(glm::value_ptr(viewport_size), "u_resolution");

    glBindVertexArray(vao_id);
    glDrawArrays(GL_LINES, 0, vertices_count);

    // draw spindle
    auto spindle_xform = glm::translate(glm::mat4(1.0f), spindle_pos);
    auto spindle_view = mvp * spindle_xform;
    shader->set_mat4(glm::value_ptr(spindle_view), "mvp");

    glBindVertexArray(spindle_vao_id);
    glDrawArrays(GL_LINES, 0, spindle_vertices_count);

    // draw bounding box
    auto bbox_size = max_pos - min_pos;

    auto bbox_translation = glm::translate(glm::mat4(1.0f), min_pos);
    auto bbox_scale = glm::scale(glm::mat4(1.0f), bbox_size);
    auto bbox_view = mvp * bbox_translation * bbox_scale;
    shader->set_mat4(glm::value_ptr(bbox_view), "mvp");

    glBindVertexArray(extents_vao_id);
    glDrawArrays(GL_LINES, 0, extents_vertices_count);

    shader->unbind();
}

void grbl::program_renderer::update(const grbl::program &pgm, const grbl::machine &cnc) {
    if (!initialized) {
        shader = new shader_program(vs_code, ps_code);
        heightmap_shader = new shader_program(vs_heightmap_code, ps_heightmap_code);

        initialize_spindle_buffers();
        initialize_program_buffers();
        initialize_extents_buffers();
        initialize_heightmap_buffers();

        initialized = true;
    }

    // update program with machine status
    // build vbo and vao
    vertices_count = update_model_vbo(pgm);

    auto machine_pos = glm::vec3{cnc.get_status().machine_pos[0],
                                 cnc.get_status().machine_pos[1],
                                 cnc.get_status().machine_pos[2]};
    auto offsets = glm::vec3{cnc.get_current_work_offset_values()[0],
                             cnc.get_current_work_offset_values()[1],
                             cnc.get_current_work_offset_values()[2]};
    spindle_pos = machine_pos - offsets;
}


void grbl::program_renderer::initialize_spindle_buffers() {
    glGenBuffers(1, &spindle_vbo_id);
    glGenVertexArrays(1, &spindle_vao_id);

    // vertex format: x, y, z, r, g, b, a
    // stride: 28 bytes
    const GLsizei size_of_vertex_in_bytes = 28;

    glBindVertexArray(spindle_vao_id);
    glBindBuffer(GL_ARRAY_BUFFER, spindle_vbo_id);

    glEnableVertexAttribArray(0); // vertices on stream 0
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, size_of_vertex_in_bytes, (void *) nullptr);

    glEnableVertexAttribArray(1); // vertex colors on stream 1
    glVertexAttribPointer(1, 4, GL_FLOAT, GL_FALSE, size_of_vertex_in_bytes, (void *) (sizeof(float) * 3));

    // unbind vao
    glBindVertexArray(0);

    // construct spindle model
    glm::vec4 col(1, 1, 0, 1);

    std::vector<float> buffer_data;

    const float spindle_length = 3.5;
    const size_t cone_granularity = 20;
    for (int i = 0; i < cone_granularity; i++) {
        float x = sinf((i / (float) cone_granularity) * 2.0f * glm::pi<float>());
        float y = cosf((i / (float) cone_granularity) * 2.0f * glm::pi<float>());

        add_line(buffer_data, {0, 0, 0}, {x, y, spindle_length}, col);
    }

    glBindBuffer(GL_ARRAY_BUFFER, spindle_vbo_id);
    glBufferData(GL_ARRAY_BUFFER, sizeof(float) * buffer_data.size(), buffer_data.data(), GL_STATIC_DRAW);

    spindle_vertices_count = buffer_data.size() * sizeof(float) / size_of_vertex_in_bytes;
}

void grbl::program_renderer::initialize_extents_buffers() {
    glGenBuffers(1, &extents_vbo_id);
    glGenVertexArrays(1, &extents_vao_id);

    // vertex format: x, y, z, r, g, b, a
    // stride: 28 bytes
    const GLsizei size_of_vertex_in_bytes = 28;

    // we're going to draw a simple box for the extents
    // box is made up of 8 vertices and 12 lines

    glBindVertexArray(extents_vao_id);
    glBindBuffer(GL_ARRAY_BUFFER, extents_vbo_id);

    glEnableVertexAttribArray(0); // vertices on stream 0
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, size_of_vertex_in_bytes, (void *) 0);

    glEnableVertexAttribArray(1); // vertex colors on stream 1
    glVertexAttribPointer(1, 4, GL_FLOAT, GL_FALSE, size_of_vertex_in_bytes, (void *) (sizeof(float) * 3));

    // unbind vao
    glBindVertexArray(0);

    // we're going to make the box as unit box and scale it when rendering
    // since we're going to reuse the same box for all loaded programs.
    // or at least we'll try to.

    // construct box
    glm::vec4 col(1, 0.7, 0.3, 1);


    // box will range from [0,0,0] to [1,1,1] and we'll use
    // translation and scaling afterward if needed to place it

    std::vector<float> buffer_data;
    // bottom plane
    add_line(buffer_data, {0, 0, 0}, {1, 0, 0}, col);
    add_line(buffer_data, {1, 0, 0}, {1, 1, 0}, col);
    add_line(buffer_data, {1, 1, 0}, {0, 1, 0}, col);
    add_line(buffer_data, {0, 1, 0}, {0, 0, 0}, col);

    // top plane
    add_line(buffer_data, {0, 0, 1}, {1, 0, 1}, col);
    add_line(buffer_data, {1, 0, 1}, {1, 1, 1}, col);
    add_line(buffer_data, {1, 1, 1}, {0, 1, 1}, col);
    add_line(buffer_data, {0, 1, 1}, {0, 0, 1}, col);

    // add vertical lines connecting the planes
    add_line(buffer_data, {0, 0, 0}, {0, 0, 1}, col);
    add_line(buffer_data, {1, 0, 0}, {1, 0, 1}, col);
    add_line(buffer_data, {1, 1, 0}, {1, 1, 1}, col);
    add_line(buffer_data, {0, 1, 0}, {0, 1, 1}, col);

    glBindBuffer(GL_ARRAY_BUFFER, extents_vbo_id);
    glBufferData(GL_ARRAY_BUFFER, sizeof(float) * buffer_data.size(), buffer_data.data(), GL_STATIC_DRAW);

    extents_vertices_count = buffer_data.size() * sizeof(float) / size_of_vertex_in_bytes;
}

void grbl::program_renderer::initialize_program_buffers() {
    glGenBuffers(1, &vbo_id);
    glGenVertexArrays(1, &vao_id);

    // vertex format: x, y, z, r, g, b, a
// stride: 28 bytes
    const GLsizei sizeOfVertexInBytes = 28;

    glBindVertexArray(vao_id);
    glBindBuffer(GL_ARRAY_BUFFER, vbo_id);

    glEnableVertexAttribArray(0); // vertices on stream 0
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, sizeOfVertexInBytes, (void *) 0);

    glEnableVertexAttribArray(1); // vertex colors on stream 1
    glVertexAttribPointer(1, 4, GL_FLOAT, GL_FALSE, sizeOfVertexInBytes, (void *) (sizeof(float) * 3));

    // unbind vao
    glBindVertexArray(0);
}

GLsizei grbl::program_renderer::update_model_vbo(const grbl::program &pgm) {
    if (!pgm.is_loaded) return 0;

    min_pos = glm::vec3(std::numeric_limits<float>::max());
    max_pos = glm::vec3(std::numeric_limits<float>::min());

    std::vector<float> buffer_data;

    grbl::grbl_parser parser;
    std::ifstream in_file{pgm.filename};
    if (in_file) {
        parser.parse(in_file);

        auto rapid_color = glm::vec4(1.0f);
        auto feed_color = glm::vec4(1.0f, 0, 0, 1);

        for (auto &c: parser.commands) {
            auto line = dynamic_cast<grbl::line_motion_cmd *>(c);
            auto arc = dynamic_cast<grbl::arc_motion_cmd *>(c);
            if (line != nullptr) {
                if (line->rapid) {
                    add_line(buffer_data, line->start, rapid_color, line->end, rapid_color);
                } else {
                    add_line(buffer_data, line->start, feed_color, line->end, feed_color);

                    update_model_extents(line->start);
                    update_model_extents(line->end);
                }


            } else if (arc != nullptr) {
                auto pieces = arc->split(0.1);
                for (auto &p: pieces) {
                    // transform arc to line
                    add_line(buffer_data, p->start, feed_color, p->end, feed_color);

                    update_model_extents(p->start);
                    update_model_extents(p->end);
                }
            }
        }
    }

    const GLsizei size_of_vertex_in_bytes = 28;
    auto number_of_vertices = buffer_data.size() * sizeof(float) / size_of_vertex_in_bytes;

    glBindBuffer(GL_ARRAY_BUFFER, vbo_id);
    glBufferData(GL_ARRAY_BUFFER, sizeof(float) * buffer_data.size(), buffer_data.data(), GL_STATIC_DRAW);

    return number_of_vertices;
}


void grbl::program_renderer::update_grid(const grbl::heightmap &grid, float exaggeration_factor) {

    glm::vec4 color = {0.5, 0.3, 0, 1};

    // this should only be called whenever the grid changes
    // therefor it should be called each and every time we
    // probe a new location. let's check that.
    std::vector<float> buffer_data;

    for (int y = 0; y < grid.y_segments; y++) {
        for (int x = 0; x < grid.x_segments; x++) {
            int current = x + y * (grid.x_segments + 1);
            int next = current + 1;
            int bottom = next + grid.x_segments;
            int bottom_next = bottom + 1;

            auto p1 = grid.vertices[current];
            auto p2 = grid.vertices[bottom];
            auto p3 = grid.vertices[next];

            // exaggerate Z
            p1.z *= exaggeration_factor;
            p2.z *= exaggeration_factor;
            p3.z *= exaggeration_factor;

            glm::vec3 normal = glm::normalize(glm::cross(p1 - p2, p3 - p1));
            add_triangle(buffer_data, p1, p2, p3, normal, color);

            p1 = grid.vertices[next];
            p2 = grid.vertices[bottom];
            p3 = grid.vertices[bottom_next];
            p1.z *= exaggeration_factor;
            p2.z *= exaggeration_factor;
            p3.z *= exaggeration_factor;

            normal = glm::normalize(glm::cross(p1 - p2, p3 - p1));
            add_triangle(buffer_data, p1, p2, p3, normal, color);
        }
    }

    const GLsizei size_of_vertex_in_bytes = 40;
    heightmap_vertices_count = buffer_data.size() * sizeof(float) / size_of_vertex_in_bytes;

    glBindBuffer(GL_ARRAY_BUFFER, heightmap_vbo_id);
    glBufferData(GL_ARRAY_BUFFER, sizeof(float) * buffer_data.size(), buffer_data.data(), GL_DYNAMIC_DRAW);
}

void grbl::program_renderer::initialize_heightmap_buffers() {
    glGenBuffers(1, &heightmap_vbo_id);
    glGenVertexArrays(1, &heightmap_vao_id);

    // for height map we are going to use solid rendering (triangles)
    // we also need to send normals so that we can better see the shape of the map through lighting

    // vertex format: v.x, v.y, v.z, n.x, n.y, n.z, r, g, b, a
    // stride: 10 * 4 bytes => 40 bytes
    const GLsizei size_of_vertex_in_bytes = 40;

    glBindVertexArray(heightmap_vao_id);
    glBindBuffer(GL_ARRAY_BUFFER, heightmap_vbo_id);

    glEnableVertexAttribArray(0); // vertices on stream 0
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, size_of_vertex_in_bytes, (void *) 0);

    glEnableVertexAttribArray(1); // normals on stream 1
    glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, size_of_vertex_in_bytes, (void *) (sizeof(float) * 3));

    glEnableVertexAttribArray(2); // vertex colors on stream 2
    glVertexAttribPointer(2, 4, GL_FLOAT, GL_FALSE, size_of_vertex_in_bytes, (void *) (sizeof(float) * 6));

    // unbind vao
    glBindVertexArray(0);
}

void grbl::program_renderer::update_model_extents(glm::vec3 point) {
    min_pos.x = std::min(min_pos.x, point.x);
    min_pos.y = std::min(min_pos.y, point.y);
    min_pos.z = std::min(min_pos.z, point.z);

    max_pos.x = std::max(max_pos.x, point.x);
    max_pos.y = std::max(max_pos.y, point.y);
    max_pos.z = std::max(max_pos.z, point.z);

}

std::string get_shader_info_log(GLuint id) {
    GLint log_length = 0;
    glGetShaderiv(id, GL_INFO_LOG_LENGTH, &log_length);

    char error_log[log_length]; // length includes the NULL character
    glGetShaderInfoLog(id, log_length, &log_length, &error_log[0]);


    return std::string(error_log, log_length);
}

std::string get_program_info_log(GLuint id) {
    GLint log_length = 0;
    glGetProgramiv(id, GL_INFO_LOG_LENGTH, &log_length);

    char error_log[log_length]; // length includes the NULL character
    glGetProgramInfoLog(id, log_length, &log_length, &error_log[0]);

    return {error_log, (size_t) log_length};
}

bool check_compile_error(GLuint shader_id) {
    GLint is_compiled = 0;
    glGetShaderiv(shader_id, GL_COMPILE_STATUS, &is_compiled);
    if (is_compiled == GL_TRUE)
        return true;

    std::cerr << "Shader compile error: " << "(id: " << shader_id << ") - " << get_shader_info_log(shader_id)
              << std::endl;
    glDeleteShader(shader_id);
    exit(1);
    return false;
}

bool check_link_error(GLuint program_id) {
    GLint is_linked = 0;
    glGetProgramiv(program_id, GL_LINK_STATUS, &is_linked);
    if (is_linked == GL_TRUE)
        return true;

    std::cerr << "Shader program link error: " << "(id: " << program_id << ") - " << get_program_info_log(program_id)
              << std::endl;
    glDeleteProgram(program_id);
    exit(1);
    return false;
}

grbl::shader_program::shader_program(const char *vs_content, const char *ps_content)
        : shader_ids{0, 0},
          program_id{0} {

    shader_ids[0] = glCreateShader(GL_VERTEX_SHADER);
    glShaderSource(shader_ids[0], 1, &vs_content, nullptr);
    glCompileShader(shader_ids[0]);
    check_compile_error(shader_ids[0]);

    shader_ids[1] = glCreateShader(GL_FRAGMENT_SHADER);
    glShaderSource(shader_ids[1], 1, &ps_content, nullptr);
    glCompileShader(shader_ids[1]);
    check_compile_error(shader_ids[1]);

    program_id = glCreateProgram();
    glAttachShader(program_id, shader_ids[0]);
    glAttachShader(program_id, shader_ids[1]);
    glLinkProgram(program_id);
    check_link_error(program_id);

    glUseProgram(program_id);
}

grbl::shader_program::~shader_program() {
    glDetachShader(program_id, shader_ids[0]);
    glDetachShader(program_id, shader_ids[1]);

    glDeleteShader(shader_ids[0]);
    glDeleteShader(shader_ids[1]);

    glDeleteProgram(program_id);
}

void grbl::shader_program::bind() const {
    glUseProgram(program_id);
}

void grbl::shader_program::unbind() {
    glUseProgram(0);
}

int grbl::shader_program::get_uniform(const char *name) {
    if (uniform_cache.count(name) == 0) {
        uniform_cache[name] = glGetUniformLocation(program_id, name);
    }
    return uniform_cache[name];
}

void grbl::shader_program::set_mat3(float *matrix, const char *uniform_name) {
    GLint location = get_uniform(uniform_name);
    if (location != -1)
        glUniformMatrix3fv(location, 1, false, matrix);
}

void grbl::shader_program::set_mat4(float *matrix, const char *uniform_name) {
    GLint location = get_uniform(uniform_name);
    if (location != -1)
        glUniformMatrix4fv(location, 1, false, matrix);
}

void grbl::shader_program::set_vec2(float *value, const char *uniform_name) {
    GLint location = get_uniform(uniform_name);
    if (location != -1)
        glUniform2fv(location, 1, value);
}

void grbl::shader_program::set_vec3(float *value, const char *uniform_name) {
    GLint location = get_uniform(uniform_name);
    if (location != -1)
        glUniform3fv(location, 1, value);
}

void grbl::shader_program::set_vec4(float *value, const char *uniform_name) {
    GLint location = get_uniform(uniform_name);
    if (location != -1)
        glUniform4fv(location, 1, value);
}

void grbl::shader_program::set_int(int value, const char *uniform_name) {
    GLint location = get_uniform(uniform_name);
    if (location != -1)
        glUniform1i(location, value);
}

void grbl::shader_program::set_float(float value, const char *uniform_name) {
    GLint location = get_uniform(uniform_name);
    if (location != -1)
        glUniform1f(location, value);
}
