#include "Gedeng/ParticleSystem.h"
#include "Gedeng/Logger.h"

#include <glm/gtc/random.hpp>

namespace Gedeng {

ParticleSystem::ParticleSystem() {
    const char *varyings[6] = {"position_out", "velocity_out", "color_out", "lifetime_out", "size_out", "type_out"};

    update_shader.add_vertex_shader("Shader/particle_update.vs");
    update_shader.add_geometry_shader("Shader/particle_update.gs");
    update_shader.add_transform_feedback(varyings, 6);
    update_shader.link();

    glGenTransformFeedbacks(1, &transform_feedback_buffer);

    render_shader.add_fragment_shader("Shader/particle_render.fs");
    render_shader.add_vertex_shader("Shader/particle_render.vs");
    render_shader.add_geometry_shader("Shader/particle_render.gs");
    render_shader.link();

    glGenQueries(1, &query);

    glGenBuffers(2, particle_buffer);
    glGenVertexArrays(2, vao);

    Particle part_initialization;
    part_initialization.type = PARTICLE_TYPE_GENERATOR;

    for (int i = 0; i < 2; i++) {
        glBindVertexArray(vao[i]);
        glBindBuffer(GL_ARRAY_BUFFER, particle_buffer[i]);
        glBufferData(GL_ARRAY_BUFFER, sizeof(Particle) * max_particle_count, NULL, GL_DYNAMIC_DRAW);
        glBufferSubData(GL_ARRAY_BUFFER, 0, sizeof(Particle), &part_initialization);

        for (int i = 0; i < 6; i++)
            glEnableVertexAttribArray(i);

        glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, sizeof(Particle), (const GLvoid *)0);  // Position
        glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, sizeof(Particle), (const GLvoid *)12); // Velocity
        glVertexAttribPointer(2, 3, GL_FLOAT, GL_FALSE, sizeof(Particle), (const GLvoid *)24); // Color
        glVertexAttribPointer(3, 1, GL_FLOAT, GL_FALSE, sizeof(Particle), (const GLvoid *)36); // Lifetime
        glVertexAttribPointer(4, 1, GL_FLOAT, GL_FALSE, sizeof(Particle), (const GLvoid *)40); // Size
        glVertexAttribPointer(5, 1, GL_FLOAT, GL_FALSE, sizeof(Particle), (const GLvoid *)44); // Type
    }

    current_read_buffer = 0;
    current_particle_count = 1;
}

void ParticleSystem::update(float delta) {
    update_shader.use();

    update_shader.setFloat("delta", delta);
    update_shader.setVec3("spawn_position", spawn_position);
    update_shader.setVec3("spawn_velocity_min", spawn_velocity_min);
    update_shader.setVec3("spawn_velocity_range", spawn_velocity_range);
    update_shader.setVec3("spawn_color", spawn_color);
    update_shader.setVec3("spawn_gravity", spawn_gravity_vector);

    update_shader.setFloat("spawn_lifetime_min", spawn_lifetime_min);
    update_shader.setFloat("spawn_lifetime_range", spawn_lifetime_range);

    update_shader.setFloat("spawn_size", spawn_size);
    update_shader.setInt("number_to_generate", 0);

    elapsed_time += delta;

    if (elapsed_time > next_generation_time) {
        update_shader.setInt("number_to_generate", target_particle_count);
        elapsed_time -= next_generation_time;

        glm::vec3 random_seed =
            glm::vec3(glm::linearRand(-10.0f, 20.0f), glm::linearRand(-10.0f, 20.0f), glm::linearRand(-10.0f, 20.0f));
        update_shader.setVec3("random_seed", random_seed);
    }

    glEnable(GL_RASTERIZER_DISCARD);
    glBindTransformFeedback(GL_TRANSFORM_FEEDBACK, transform_feedback_buffer);

    glBindVertexArray(vao[current_read_buffer]);
    glEnableVertexAttribArray(1); // Re-enable velocity

    glBindBufferBase(GL_TRANSFORM_FEEDBACK_BUFFER, 0, particle_buffer[1 - current_read_buffer]);

    glBeginQuery(GL_TRANSFORM_FEEDBACK_PRIMITIVES_WRITTEN, query);
    glBeginTransformFeedback(GL_POINTS);

    glDrawArrays(GL_POINTS, 0, current_particle_count);

    glEndTransformFeedback();

    glEndQuery(GL_TRANSFORM_FEEDBACK_PRIMITIVES_WRITTEN);
    glGetQueryObjectiv(query, GL_QUERY_RESULT, &current_particle_count);

    current_read_buffer = 1 - current_read_buffer;

    glBindTransformFeedback(GL_TRANSFORM_FEEDBACK, 0);
}

void ParticleSystem::set_camera(const Camera &camera) {
    projection_matrix = camera.get_projection();
    view_matrix = camera.get_view();

    quad1 = glm::cross(camera.forward(), camera.up());
    quad1 = glm::normalize(quad1);

    quad2 = glm::cross(camera.forward(), quad1);
    quad2 = glm::normalize(quad2);
}

void ParticleSystem::render() {
    glEnable(GL_BLEND);
    glBlendFunc(GL_SRC_ALPHA, GL_ONE);
    glDepthMask(0);

    glDisable(GL_RASTERIZER_DISCARD);
    render_shader.use();
    render_shader.setMat4("projection", projection_matrix);
    render_shader.setMat4("view", view_matrix);
    render_shader.setVec3("quad1", quad1);
    render_shader.setVec3("quad2", quad2);

    texture->bind_to(0);

    glBindVertexArray(vao[current_read_buffer]);
    glDisableVertexAttribArray(1); // Disable velocity, because we don't need it for rendering

    glDrawArrays(GL_POINTS, 0, current_particle_count);

    glDepthMask(1);
    glDisable(GL_BLEND);
}

void ParticleSystem::set_position(glm::vec3 position) {
    this->spawn_position = position;
}
void ParticleSystem::set_velocity(glm::vec3 min, glm::vec3 max) {
    this->spawn_velocity_min = min;
    this->spawn_velocity_range = max - min;
}
void ParticleSystem::set_gravity(glm::vec3 gravity) {
    this->spawn_gravity_vector = gravity;
}
void ParticleSystem::set_color(glm::vec3 color) {
    this->spawn_color = color;
}
void ParticleSystem::set_lifetime(float min, float max) {
    this->spawn_lifetime_min = min;
    this->spawn_lifetime_range = max - min;
}
void ParticleSystem::set_size(float size) {
    this->spawn_size = size;
}
void ParticleSystem::set_interval(float interval) {
    this->next_generation_time = interval;
}
void ParticleSystem::set_number_of_particles(int number) {
    this->target_particle_count = number;
}
void ParticleSystem::set_texture(Texture *texture) {
    this->texture = texture;
}

} // namespace Gedeng