#pragma once

#include <cgl/buffer.hpp>
#include <cgl/shaderprogram.hpp>
#include <cgl/vertexarray.hpp>
#include <cgl/pipeline.hpp>

#include <tc/groups.hpp>

#include <Eigen/Eigen>

#include <entt/entt.hpp>

#include "mirror.hpp"
#include "geometry.hpp"
#include "solver.hpp"

#include <shaders.hpp>

namespace vis {
    template<int R_, int D_, int G_>
    struct Structure {
        static constexpr auto Rank = R_;
        static constexpr auto Dim = D_;
        static constexpr auto Grade = G_;

        using Affine = Eigen::Transform<float, Dim, Eigen::Affine>;

        using Vertex = Eigen::Vector<float, Dim>;
        using Color = Eigen::Vector<float, 3>;
        using Cell = Eigen::Array<unsigned int, Grade, 1>;

        Points points;
        Hull<Grade> hull;

        std::vector<char> enabled;
        std::vector<Eigen::Vector3f> colors;

        Affine transform = Affine::Identity();

        template<typename P, typename H>
        explicit Structure(P &&points_, H &&hull_, Color color_ = Color::Ones()):
            points(std::forward<P>(points_)),
            hull(std::forward<H>(hull_)),
            enabled(hull.tilings.size(), true),
            colors(hull.tilings.size(), color_),
            transform(Affine::Identity()) {
        }
    };

    template<typename Str_>
    struct VBOs {
        using Str = Str_;

        struct Uniform {
            Eigen::Matrix4f linear;
            Eigen::Vector4f translation;
        };

        struct Command {
            unsigned int count, instanceCount, first, baseInstance;
        };

        cgl::Buffer<typename Str::Vertex> vertices;
        cgl::Buffer<typename Str::Color> colors;
        cgl::Buffer<typename Str::Cell> indices;
        cgl::Buffer<Uniform> uniform;
        cgl::Buffer<Command> commands;
    };

    template<typename Str>
    void upload_structure(entt::registry &registry) {
        auto view = registry.view<Str, VBOs<Str>>();

        for (auto [entity, structure, vbos]: view.each()) {
            auto vertices = structure.points.verts.colwise();
            auto indices = structure.hull.inds.colwise();

            vbos.vertices.put(vertices.begin(), vertices.end());
            vbos.indices.put(indices.begin(), indices.end());
        }
    }

    template<typename Str>
    void upload_uniforms(entt::registry &registry) {
        auto view = registry.view<Str, VBOs<Str>>();

        for (auto [entity, structure, vbos]: view.each()) {
            auto colors = structure.colors;
            typename VBOs<Str>::Uniform uniform{
                structure.transform.linear(),
                structure.transform.translation(),
            };

            vbos.colors.put(colors.begin(), colors.end());
            vbos.uniform.put(uniform, GL_STREAM_DRAW);
        }
    }

    template<typename Str>
    void upload_commands(entt::registry &registry) {
        auto view = registry.view<Str, VBOs<Str>>();

        for (auto [entity, structure, vbos]: view.each()) {
            const auto &tilings = structure.hull.tilings;

            std::vector<typename VBOs<Str>::Command> commands;

            for (unsigned int i = 0; i < tilings.size(); ++i) {
                if (structure.enabled[i]) {
                    auto [first, count] = tilings[i];
                    commands.push_back({(unsigned int) count, 1, (unsigned int) first, i});
                }
            }

            vbos.commands.put(commands.begin(), commands.end(), GL_STREAM_DRAW);
        }
    }

    template<typename Str_>
    struct SliceRenderer {
        using Str = Str_;

        cgl::pgm::vert defer = cgl::pgm::vert(shaders::deferred_vs_glsl);
        cgl::pgm::geom slice = cgl::pgm::geom(shaders::slice_gm_glsl);
        cgl::pgm::frag solid = cgl::pgm::frag(shaders::solid_fs_glsl);

        cgl::pipeline pipe;

        cgl::VertexArray vao;

        SliceRenderer() {
            pipe.stage(defer);
            pipe.stage(slice);
            pipe.stage(solid);

            vao.iformat(0, 4, GL_UNSIGNED_INT);
            vao.format(1, 3, GL_FLOAT);

            glVertexArrayBindingDivisor(vao, 1, 1);
        }

        void operator()(entt::registry &reg) {
            auto view = reg.view<VBOs<Str>>();

            for (auto [entity, vbos]: view.each()) {
                glBindProgramPipeline(pipe);

                glBindBufferBase(GL_UNIFORM_BUFFER, 2, vbos.uniform);
                glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, vbos.vertices);

                vao.vertexBuffer(0, vbos.indices);
                vao.vertexBuffer(1, vbos.colors);

                glBindBuffer(GL_DRAW_INDIRECT_BUFFER, vbos.commands);

                glBindVertexArray(vao);
                glMultiDrawArraysIndirect(GL_POINTS, nullptr, vbos.commands.count(), 0);
                glBindVertexArray(0);

                glBindProgramPipeline(0);
            }
        }
    };
}