#include <gfx/render.hpp>

#include <algo/matrix.hpp>
#include <gfx/camera.hpp>
#include <gfx/constants.hpp>
#include <gfx/mesh.hpp>
#include <gfx/shader.hpp>
#include <osi/opengl.hpp>
#include <utils/assumptions.hpp>
#include <utils/math.hpp>

#include <iostream>
#include <vector>

OpenGLContext::OpenGLContext()
{
    glClearColor(0.2f, 0.3f, 0.3f, 1.0f);
    glBlendFunc(GL_ONE, GL_ONE_MINUS_SRC_ALPHA);
    glEnable(GL_MULTISAMPLE);

    glEnable(GL_LINE_SMOOTH);
    glLineWidth(2.0f);

    // GLfloat aliasWidthRange[2] = {0.0f, 0.0f};
    // glGetFloatv(GL_ALIASED_LINE_WIDTH_RANGE, aliasWidthRange);
    // GLfloat smoothWidthRange[2] = {0.0f, 0.0f};
    // glGetFloatv(GL_SMOOTH_LINE_WIDTH_RANGE, smoothWidthRange);

    glCreateBuffers(1, &camera_buffer);
    glNamedBufferStorage(camera_buffer, sizeof(CameraUBO), 0, GL_DYNAMIC_STORAGE_BIT);
    ASSUMPTION(glGetError() == GL_NO_ERROR);

    glCreateBuffers(1, &material_buffer);
    glNamedBufferStorage(material_buffer, sizeof(Material), 0, GL_DYNAMIC_STORAGE_BIT);
    ASSUMPTION(glGetError() == GL_NO_ERROR);

    glCreateBuffers(1, &lights_buffer);
    glNamedBufferStorage(lights_buffer, lights_buffer_size * sizeof(LightUBO), 0, GL_DYNAMIC_STORAGE_BIT);
    ASSUMPTION(glGetError() == GL_NO_ERROR);
}

OpenGLContext::~OpenGLContext()
{
    glDeleteBuffers(1, &camera_buffer);
    glDeleteBuffers(1, &material_buffer);
    glDeleteBuffers(1, &lights_buffer);
}

std::vector<std::pair<std::string, shader_ptr>> sort_shaders(const std::map<std::string, shader_ptr> &my_shaders)
{
    std::vector<std::pair<std::string, shader_ptr>> sorted_shaders;
    std::transform(
        my_shaders.begin(), my_shaders.end(),
        std::back_inserter(sorted_shaders),
        [](const auto& pair) { return pair; }
    );

    std::sort(sorted_shaders.begin(), sorted_shaders.end(),
        [](const std::pair<std::string, shader_ptr>& a, const std::pair<std::string, shader_ptr>& b) {
            return a.second->getType() < b.second->getType();
        }
    );

    return sorted_shaders;
}

mesh_ptr get_skybox_mesh(node_ptr skybox)
{
    for (auto object : skybox->getObjects())
        if (auto mesh = std::dynamic_pointer_cast<Mesh>(object); mesh)
            return mesh;

    return nullptr;
}

void setup_skybox(OpenGLContext &gfx_context, node_ptr skybox)
{
    if (skybox)
    {
        auto mesh = get_skybox_mesh(skybox);
        if (mesh)
            gfx_context.skybox_texture = mesh->getTexture()->getID();
    }
}

bool opaqueness_check(mesh_ptr mesh, bool opaque_render)
{
    if (mesh->getMaterial().diffuse_color.w < 1.0f && opaque_render)
        return false;
    if (mesh->getMaterial().diffuse_color.w >= 1.0f && !opaque_render)
        return false;

    return true;
}

void send_buffer_data(GLuint buffer_name, GLsizei size, const void *data)
{
    glNamedBufferSubData(buffer_name, 0, size, data);
    ASSUMPTION(glGetError() == GL_NO_ERROR);
}

void resize_buffer(GLuint buffer_name, GLsizei size)
{
    glDeleteBuffers(1, &buffer_name);
    ASSUMPTION(glGetError() == GL_NO_ERROR);
    glCreateBuffers(1, &buffer_name);
    ASSUMPTION(glGetError() == GL_NO_ERROR);
    glNamedBufferStorage(buffer_name, size, 0, GL_DYNAMIC_STORAGE_BIT);
    ASSUMPTION(glGetError() == GL_NO_ERROR);
}

void setup_camera_buffer(OpenGLContext &gfx_context, camera_ptr camera)
{
    gfx_context.camera_ubo.view = get_view_matrix(camera);
    gfx_context.camera_ubo.projection = get_projection_matrix(camera);
    gfx_context.camera_ubo.position = glm::vec4(camera->getNode()->getPositionWorld(), 1.0f);
}

void setup_lights_buffer(OpenGLContext &gfx_context, std::vector<light_ptr> &lights)
{
    if (static_cast<GLsizeiptr>(lights.size()) > gfx_context.lights_buffer_size)
    {
        gfx_context.lights_buffer_size = lights.size();
        resize_buffer(gfx_context.lights_buffer, static_cast<GLsizei>(gfx_context.lights_buffer_size) * sizeof(LightUBO));
    }

    gfx_context.lights_ubos.clear();

    for (auto &light : lights)
        gfx_context.lights_ubos.emplace_back(light->getPosition(), 
                                             light->getAmbientColor(), 
                                             light->getDiffuseColor(), 
                                             light->getSpecularColor(), 
                                             light->getConeDirection());
}

void setup_buffers(OpenGLContext &gfx_context, camera_ptr camera, std::vector<light_ptr> &lights)
{
    setup_camera_buffer(gfx_context, camera);
    send_buffer_data(gfx_context.camera_buffer, sizeof(CameraUBO), &gfx_context.camera_ubo);

    setup_lights_buffer(gfx_context, lights);
    send_buffer_data(gfx_context.lights_buffer, static_cast<GLsizei>(gfx_context.lights_buffer_size) * sizeof(LightUBO), gfx_context.lights_ubos.data());
}

void setup_shader_basic(shader_ptr shader_program, const OpenGLContext &gfx_context)
{
    glEnable(GL_DEPTH_TEST);
    // glEnable(GL_CULL_FACE);

    shader_program->setMat4(SHD_VIEW_MAT, gfx_context.camera_ubo.view);
    shader_program->setMat4(SHD_PROJ_MAT, gfx_context.camera_ubo.projection);
}

void setup_shader_lit(shader_ptr shader_program, const OpenGLContext &gfx_context)
{
    glEnable(GL_DEPTH_TEST);
    // glEnable(GL_CULL_FACE);

    shader_program->setUBO(SHD_CAMERA, gfx_context.camera_buffer);
    shader_program->setSSBO(SHD_LIGHTS, gfx_context.lights_buffer);
    shader_program->setTU(SHD_SKYBOX_TEX, gfx_context.skybox_texture);
}

void setup_shader_skybox(shader_ptr shader_program, const OpenGLContext &gfx_context)
{
    glDisable(GL_DEPTH_TEST);
    glDisable(GL_CULL_FACE);

    shader_program->setUBO(SHD_CAMERA, gfx_context.camera_buffer);
}

void setup_shader_interface(shader_ptr shader_program, const OpenGLContext &gfx_context)
{
    glDisable(GL_DEPTH_TEST);
    // glEnable(GL_CULL_FACE);

    shader_program->setMat4(SHD_VIEW_MAT, gfx_context.camera_ubo.view);
    shader_program->setMat4(SHD_PROJ_MAT, gfx_context.camera_ubo.projection);
}

void setup_shader(std::map<std::string, shader_ptr> &my_shaders, const std::string &name, const OpenGLContext &gfx_context)
{
    ASSUMPTION(my_shaders.contains(name));

    if (name == "basic")
        setup_shader_basic(my_shaders["basic"], gfx_context);
    else if (name == "lit")
        setup_shader_lit(my_shaders["lit"], gfx_context);
    else if (name == "skybox")
        setup_shader_skybox(my_shaders["skybox"], gfx_context);
    else if (name == "interface")
        setup_shader_interface(my_shaders["interface"], gfx_context);
}

void send_mesh_data_basic_shader(const shader_ptr &shader_program, const OpenGLContext &gfx_context, mesh_ptr mesh)
{
    shader_program->setMat4(SHD_MODEL_MAT, gfx_context.model_mat_holder);
    shader_program->setVec4(SHD_OBJECT_COLOR, mesh->getColor());
}

void send_mesh_data_lit_shader(const shader_ptr &shader_program, const OpenGLContext &gfx_context, mesh_ptr mesh)
{
    shader_program->setMat4(SHD_MODEL_MAT, gfx_context.model_mat_holder);
    send_buffer_data(gfx_context.material_buffer, sizeof(Material), &mesh->getMaterial());
    shader_program->setUBO(SHD_OBJECT_MATERIAL, gfx_context.material_buffer);
    shader_program->setBool(SHD_TEXTURE_BOOL, mesh->hasTexture());
    shader_program->setTU(SHD_ALBEDO_TEX, mesh->getTexture()->getID());
}

void send_mesh_data_skybox_shader(const shader_ptr &shader_program, const OpenGLContext &gfx_context, mesh_ptr mesh)
{
    shader_program->setTU(SHD_SKYBOX_TEX, mesh->getTexture()->getID());
}

void send_mesh_data_interface_shader(const shader_ptr &shader_program, const OpenGLContext &gfx_context, mesh_ptr mesh)
{
    shader_program->setMat4(SHD_MODEL_MAT, gfx_context.model_mat_holder);
    shader_program->setVec4(SHD_OBJECT_COLOR, mesh->getColor());
}

void send_mesh_data(const std::pair<std::string, shader_ptr> &shader, const OpenGLContext &gfx_context, mesh_ptr mesh)
{
    const auto &[shader_name, shader_program] = shader;
    if (shader_name == "basic")
        send_mesh_data_basic_shader(shader_program, gfx_context, mesh);
    else if (shader_name == "lit")
        send_mesh_data_lit_shader(shader_program, gfx_context, mesh);
    else if (shader_name == "skybox")
        send_mesh_data_skybox_shader(shader_program, gfx_context, mesh);
    else if (shader_name == "interface")
        send_mesh_data_interface_shader(shader_program, gfx_context, mesh);
}

void draw_objects(const std::pair<std::string, shader_ptr> &shader, std::vector<light_ptr> lights,
                  camera_ptr camera,
                  const std::vector<node_ptr> &scene,
                  const OpenGLContext &gfx_context)
{
    auto &[shader_name, shader_program] = shader;

    for (auto &node : scene)
    {
        for (auto &object : node->getObjects())
        {
            auto mesh = std::dynamic_pointer_cast<Mesh>(object);
            if (!mesh || shader_name != mesh->getShaderType() || !mesh->isVisible()
                || !opaqueness_check(mesh, gfx_context.opaque_render))
                continue;

            gfx_context.model_mat_holder = node->getModelMatWorld();
            
            send_mesh_data(shader, gfx_context, mesh);

            mesh->draw();
        }

        draw_objects(shader, lights, camera, node->getChildren(), gfx_context);
    }
}

void gfx_draw(std::map<std::string, shader_ptr> &my_shaders, 
              std::vector<light_ptr> &lights,
              camera_ptr camera,
              node_ptr skybox,
              const std::vector<node_ptr> &scene,
              OpenGLContext &gfx_context)
{
    glDepthMask(GL_TRUE);

    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT | GL_STENCIL_BUFFER_BIT);
    ASSUMPTION(glGetError() == GL_NO_ERROR);

    glViewport(0, 0, static_cast<GLsizei>(camera->getWindowSize().x), static_cast<GLsizei>(camera->getWindowSize().y));
    ASSUMPTION(glGetError() == GL_NO_ERROR);

    setup_skybox(gfx_context, skybox);
    setup_buffers(gfx_context, camera, lights);

    auto sorted_shaders = sort_shaders(my_shaders);

    /* OPAQUE RENDER */
    gfx_context.opaque_render = true;

    glDisable(GL_BLEND);

    for (auto [shader_name, shader_program] : sorted_shaders)
    {
        shader_program->use();
        setup_shader(my_shaders, shader_name, gfx_context);
        draw_objects(std::pair(shader_name, shader_program), lights, camera, scene, gfx_context);
    }

    /* TRANSPARENT RENDER */
    gfx_context.opaque_render = false;

    glDepthMask(GL_FALSE);
    glEnable(GL_BLEND);

    for (auto [shader_name, shader_program] : sorted_shaders)
    {
        shader_program->use();
        setup_shader(my_shaders, shader_name, gfx_context);
        draw_objects(std::pair(shader_name, shader_program), lights, camera, scene, gfx_context);
    }
}
