#pragma once
#ifndef _NB_BUFFER
#define _NB_BUFFER

#include <GLLoad.h>

#include <exception>
#include <map>
#include <string>
#include <type_traits>
#include <vector>

#include "Shader.h"

#define THROW_BUFFER_ERROR(msg) throw BufferError(msg, __FILE__, __LINE__);
#define THROW_VAO_ERROR(msg) throw VAOError(msg, __FILE__, __LINE__);

namespace NB{

typedef std::vector<unsigned char> RawVec;

struct VertexAttributePointer {
    GLuint buffer = 0;
    int32_t offset = 0;
    GLsizei stride = -1;
    GLuint divisor = 0;
};

struct VertexAttribute {
    GLint GLSLSize;
    GLenum GLSLType;
    GLboolean GLSLNormalization;
    VertexAttributePointer ptr;
};

typedef std::vector<VertexAttribute> VertexAttributeList;

template<typename T>
RawVec vectorToRaw(const std::vector<T>& vec) {
    unsigned int num_bytes = vec.size() * sizeof(T);
    RawVec ret(num_bytes);
    memcpy(ret.data(), vec.data(), num_bytes);
    return ret;
}

template<typename T, typename S>
RawVec concatVectorsToRaw(const std::vector<T>& vec1, const std::vector<S>& vec2) {
    RawVec vec1_raw = vectorToRaw<T>(vec1);
    RawVec vec2_raw = vectorToRaw<S>(vec2);
    unsigned int vec1_raw_size = vec1_raw.size();
    unsigned int vec2_raw_size = vec2_raw.size();
    RawVec ret(vec1_raw_size + vec2_raw_size);
    memcpy(ret.data(), vec1_raw.data(), vec1_raw_size);
    memcpy(ret.data() + vec1_raw_size, vec2_raw.data(), vec2_raw_size);
    return ret;
}

static uint8_t GLSLTypeSize(GLenum type) {
    switch(type) {
    case GL_SHORT:
    case GL_UNSIGNED_SHORT:
    case GL_HALF_FLOAT:
        return 2;
        break;
    case GL_INT:
    case GL_UNSIGNED_INT:
    case GL_FLOAT:
        return 4;
        break;
    case GL_DOUBLE:
        return 8;
        break;
    default:
        return 1;
        break;
    }
}

template<typename T>
std::vector<T> rawToVector(const RawVec& vec) {
    if (vec.size() % sizeof(T) != 0) {
        throw std::runtime_error("Data size does not align to std::vector<" + std::string(typeid(T).name()) + ">.");
    }
    unsigned int num_elmts = vec.size() / sizeof(T);
    std::vector<T> ret(num_elmts);
    memcpy(ret.data(), vec.data(), vec.size());

    return ret;
}

class BufferError : public std::runtime_error {
public:
    const bool error;
    BufferError(const std::string& msg, const std::string& file="", int line=-1) 
        : std::runtime_error(NB::formatDebugString(msg, file, line)), error(true) {}
    BufferError(bool isError=true) : std::runtime_error(""), error(isError) {}
};

class VAOError : public std::runtime_error {
public:
    const bool error;
    VAOError(const std::string& msg, const std::string& file="", int line=-1) 
        : std::runtime_error(NB::formatDebugString(msg, file, line)), error(true) {}
    VAOError(bool isError=true) : std::runtime_error(""), error(isError) {}
};

template<GLenum BT, typename T=unsigned char>
class BufferManager {
public:
    const GLenum BufferType;

    BufferManager(
        bool keep_local=true,
        GLenum usage=GL_STATIC_DRAW
    ) : BufferType(BT), _keep_local(keep_local), _usage(usage), _size(0) {
            static_assert( !(BT==GL_ARRAY_BUFFER) || std::is_same<T, unsigned char>::value,
                            "'GL_ARRAY_BUFFER' buffer manager must have type of 'unsigned char'." );
            static_assert( !(BT==GL_ELEMENT_ARRAY_BUFFER) || std::is_integral<T>::value,
                            "'GL_ELEMENT_ARRAY_BUFFER' buffer manager must be integral type." );
            glGenBuffers(1, &_id);
        }
    BufferManager(
        const std::vector<T>& data,
        bool keep_local=true,
        GLenum usage=GL_STATIC_DRAW
    ) : BufferManager<BT, T>(keep_local, usage) { setData(data); }
    BufferManager(
        unsigned int size,
        bool keep_local=true,
        GLenum usage=GL_STATIC_DRAW
    ) : BufferManager<BT, T>(std::vector<T>(size), keep_local, usage) {}
    ~BufferManager() {
        glDeleteBuffers(1, &_id);
    }

    GLenum usageType() const { return _usage; }
    unsigned int id() const { return _id; }
    GLsizei capacity() const { return _size; }
    bool isInitialized() const { return glIsBuffer(_id); }
    bool keepLocal() const { return _keep_local; }
    std::vector<T> getDataDirect() const {
        std::vector<T> ret(_size);
        bind();
        glGetBufferSubData(BT, 0, sizeof(T)*_size, ret.data());

        return ret;
    }
    std::vector<T> getData() const {
        if(_keep_local) {
            return _data;
        } else {
            return getDataDirect();
        }
    }
    GLuint bind() const { glBindBuffer(BT, _id); return _id; }
    void unbind() const { glBindBuffer(BT, 0); }

    GLenum usageType(GLenum usage) { _usage=usage; return usageType(); }
    bool keepLocal(bool keep_local) {
        if(keep_local && !_keep_local) {
            _data = getDataDirect();
        }
        if (!keep_local && _keep_local && _data.size() != 0) {
            _data.erase(_data.begin(), _data.end());
        }
        _keep_local = keep_local;
        return _keep_local;
    }
    void setData(const std::vector<T>& set_data) {
        if (!glfwGetCurrentContext()) {
            THROW_BUFFER_ERROR("No OpenGL context.");
        }
        _size = set_data.size();
        bind();
        glBufferData(BT, _size*sizeof(T), set_data.data(), _usage);
        if (_keep_local) { _data = set_data; }
    }
    unsigned int setSubData(const std::vector<T>& newData, unsigned int offset=0) {
        if (newData.size()+offset > _size) {
            THROW_BUFFER_ERROR(
                "Attempting to overflow buffer of capacity " + std::to_string(_size) + 
                " with data type size " + std::to_string(sizeof(T)) + ".");
        }
        bind();
        glBufferSubData(BT, offset*sizeof(T), sizeof(T)*newData.size(), newData.data());
        if(_keep_local) {
            memcpy(_data.data()+offset*sizeof(T), newData.data(), sizeof(T)*newData.size());
        }
        return _size;


    }
    unsigned int resize(unsigned int n) {
        std::vector<T> new_vector = getData();
        new_vector.resize(n);
        setData(new_vector);
        return n;
    }

private:
    GLenum _usage;
    GLuint _id = 0;
    std::vector<T> _data;
    unsigned int _size;
    bool _keep_local;
};

template<typename ElmtIdxType=unsigned int>
class VAOManager {
public:
    VAOManager() : _elmt_buffer(nullptr){
        glGenVertexArrays(1, &_id);
    }
    VAOManager(
        std::vector<BufferManager<GL_ARRAY_BUFFER>*> vert_bufs,
        BufferManager<GL_ELEMENT_ARRAY_BUFFER, ElmtIdxType>* elmt_buf = nullptr,
        const VertexAttributeList& vert_attrs = {}
    ) : VAOManager() {
        _elmt_buffer = elmt_buf;
        for (BufferManager<GL_ARRAY_BUFFER>* vb : vert_bufs) {
            _vert_buffers[vb->id()] = vb;
        }
        if (vert_attrs.size()!=0) {
            generate(vert_attrs);
        }
    }
    VAOManager(
        BufferManager<GL_ARRAY_BUFFER>* vert_bufs,
        BufferManager<GL_ELEMENT_ARRAY_BUFFER, ElmtIdxType>* elmt_buf = nullptr,
        const VertexAttributeList& vert_attrs = {}
    ) : VAOManager(std::vector<BufferManager<GL_ARRAY_BUFFER>*>(1, vert_bufs), elmt_buf, vert_attrs) {}

    VertexAttributeList getLayout() const { return _vert_attrs; }
    GLuint id() const { return _id; }
    const VertexAttribute& operator[](unsigned int i) { return _vert_attrs[i]; }
    unsigned int vertSize(unsigned int i) const {
        // PROBLEM IS ALL VAPs POINT TO IBO rn
        unsigned int vert_size = 0;
        for (const VertexAttribute& va : _vert_attrs) {
            if (va.ptr.buffer == i) {
                vert_size += va.GLSLSize * GLSLTypeSize(va.GLSLType);
            }
        }
        return vert_size;
    }
    std::vector<GLuint> getVBOs() {
        std::vector<GLuint> ret;
        for(std::map<GLuint, BufferManager<GL_ARRAY_BUFFER>*>::const_iterator it = _vert_buffers.begin(); it != _vert_buffers.end(); ++it) {
            ret.push_back(it->first);
        }
        return ret;
    }

    void bind() const {
        glBindVertexArray(_id); 
        if (_elmt_buffer) { _elmt_buffer->bind(); }    
    }
    void unbind() const {
        glBindVertexArray(0);
        if (_elmt_buffer) { _elmt_buffer->bind(); }
    }
    void addVBO(BufferManager<GL_ARRAY_BUFFER>* vert_buf) {
        GLuint vert_id = vert_buf->id();
        if (_vert_buffers.find(vert_id) == _vert_buffers.end()) {
            _vert_buffers[vert_buf->id()] = vert_buf;
        } else {
            THROW_VAO_ERROR("Attempting to add identical VBO id of " + std::to_string(vert_id) + ".");
        }
    }
    void changeEBO(BufferManager<GL_ELEMENT_ARRAY_BUFFER, ElmtIdxType>* ebo) {
        _elmt_buffer = ebo;
    }
    VertexAttributeList addVertexAttributes(const VertexAttributeList& vert_attrs) { 
        try {
            throw checkValid(vert_attrs);
        } catch (VAOError vaoe) {
            if (vaoe.error) {
                THROW_VAO_ERROR(vaoe.what());
            }
        }
        unsigned int num_attrs_old = _vert_attrs.size();
        unsigned int num_attrs_new = num_attrs_old+vert_attrs.size();
        const VertexAttribute* va;
        bind();
        for (int i = num_attrs_old; i < num_attrs_new; ++i) {
            va = &(vert_attrs.data()[i-num_attrs_old]);
            _vert_attrs.emplace_back(*va);
            _vert_buffers[va->ptr.buffer]->bind();
            glVertexAttribPointer(
                i,
                va->GLSLSize,
                va->GLSLType,
                va->GLSLNormalization,
                va->ptr.stride,
                (void*)va->ptr.offset
            );
            glVertexAttribDivisor(i, va->ptr.divisor);
            glEnableVertexAttribArray(i);
            glVertexAttribDivisor(i, va->ptr.divisor);

        }
        unbind();
        return _vert_attrs;
    }
    VertexAttributeList addVertexAttributes(VertexAttribute vert_attr) { return addVertexAttributes({vert_attr}); }
    VertexAttributeList generate() {
        bind();
        unsigned int num_attrs = _vert_attrs.size();
        VertexAttribute* va;
        for (int i = 0; i < num_attrs; ++i) {
            va = &(_vert_attrs[i]);
            _vert_buffers[va->ptr.buffer]->bind();
            glVertexAttribPointer(
                i,
                va->GLSLSize,
                va->GLSLType,
                va->GLSLNormalization,
                va->ptr.stride,
                (void*)va->ptr.offset
            );
            glVertexAttribDivisor(i, va->ptr.divisor);
            glEnableVertexAttribArray(i);
        }
        if (_elmt_buffer != nullptr) {
            _elmt_buffer->bind();
        }
        unbind();
        return _vert_attrs;
    }
    VertexAttributeList generate(VertexAttributeList vert_attrs) {
        try {
            throw checkValid(vert_attrs);
        } catch (VAOError vaoe) {
            if (vaoe.error) {
                THROW_VAO_ERROR(vaoe.what());
            }
        }
        _vert_attrs.swap(vert_attrs);
        return generate();
    }
    VertexAttributeList changeLayout(unsigned int i, VertexAttributePointer vap) {
        _vert_attrs[i].ptr = vap;
        bind();
        const VertexAttribute* va = &(_vert_attrs[i]);
        glDisableVertexAttribArray(i);
        _vert_buffers[va->ptr.buffer]->bind();
        glVertexAttribPointer(
            i,
            va->GLSLSize,
            va->GLSLType,
            va->GLSLNormalization,
            va->ptr.stride,
            (void*)va->ptr.offset
        );
        glVertexAttribDivisor(i, va->ptr.divisor);
        glEnableVertexAttribArray(i);
        unbind();
        return _vert_attrs;
    }

    VAOError checkValid(const VertexAttributeList& vert_attrs) {
        GLuint va_id;
        unsigned int num_attrs = vert_attrs.size();
        for (int i = 0; i < num_attrs; ++i) {
            va_id = vert_attrs[i].ptr.buffer;
            if (_vert_buffers.find(va_id) == _vert_buffers.cend()) {
                return VAOError("Attempting to point to unknown VBO of id " + std::to_string(va_id) + 
                " at Vertex Attribute " + std::to_string(i) + ".");
            }
            if (vert_attrs[i].ptr.stride<0) {
                return VAOError("Invalid stride value at Vertex Attribute " + std::to_string(i) + ".");
            }
        }
        return VAOError(false);
    }

private:
    GLuint _id;
    std::map<GLuint, BufferManager<GL_ARRAY_BUFFER>*> _vert_buffers;
    BufferManager<GL_ELEMENT_ARRAY_BUFFER, ElmtIdxType>* _elmt_buffer;
    VertexAttributeList _vert_attrs;
};

}

#endif