#ifndef _BASIC_MODEL_HPP_
#define _BASIC_MODEL_HPP_

#include <cmath>

#include "tmath.hpp"
#include "rusty.hpp"

struct Vertex {
    tmath::fvec3 position;
    tmath::fvec3 normal;
    tmath::fvec2 texcoord;
    // 0 - attribute 1 - texture
    tmath::ivec2 index;
};

struct VertexAttribute {
    tmath::fmat4 matrix = tmath::fmat4::identity();
    tmath::fmat4 nmatrix = tmath::fmat4::identity();
    tmath::fvec4 color = {1.0f, 1.0f, 1.0f, 1.0f};
    tmath::fvec4 mixcolor = {0.0f, 0.0f, 0.0f, 0.0f};
};

struct ModelUnit {
    uint32_t vert_offset;
    uint32_t vert_count;
    uint32_t index_offset;
    uint32_t index_count;
};

class VertCollection {
public:
    ModelUnit gen_square(float tx = 0.0f, float ty = 0.0f, float tw = 1.0f, float th = 1.0f) {
        uint32_t offset = out_v.size();
        uint32_t ioffset = out_i.size();
        // vert_count = 4;
        out_v.push_back({{-1.0f, -1.0f, 0.0f}, {0.0f, 0.0f, 1.0f}, {tx, ty}});
        out_v.push_back({{-1.0f, 1.0f, 0.0f}, {0.0f, 0.0f, 1.0f}, {tx, ty + th}});
        out_v.push_back({{1.0f, 1.0f, 0.0f}, {0.0f, 0.0f, 1.0f}, {tx + tw, ty + th}});
        out_v.push_back({{1.0f, -1.0f, 0.0f}, {0.0f, 0.0f, 1.0f}, {tx + tw, ty}});
        out_i.insert(out_i.end(), {0 + offset, 1 + offset, 2 + offset, 2 + offset, 3 + offset, 0 + offset});
        return ModelUnit{offset, (uint32_t)out_v.size() - offset, ioffset, (uint32_t)out_i.size() - ioffset};
    }

    ModelUnit gen_texture(uint32_t tilew, uint32_t tileh, float posx, float posy, float texw, float texh) {
        uint32_t offset = out_v.size();
        uint32_t ioffset = out_i.size();
        float tx = 1.0f / tilew;
        float ty = 1.0f / tileh;
        float mx = texw * tx;
        float my = texh * ty;
        // vert_count = (tilew + 1) * (tileyh + 1);
        for(uint32_t y = 0; y <= tileh; ++y) {
            for(uint32_t x = 0; x <= tilew; ++x)
                out_v.push_back({{posx + x * mx, posy + y * my, 0.0f}, {0.0f, 0.0f, 1.0f}, {x * tx, y * ty}});
        }
        for(uint32_t y = 0; y < tileh; ++y) {
            for(uint32_t x = 0; x < tilew; ++x) {
                uint32_t idx = y * (tileh + 1) + x + offset;
                out_i.insert(out_i.end(), {idx, idx + tileh + 1, idx + 1, idx  + 1, idx + tileh + 1, idx + tileh + 2});
            }
        }
        return ModelUnit{offset, (uint32_t)out_v.size() - offset, ioffset, (uint32_t)out_i.size() - ioffset};
    }

    ModelUnit gen_cube() {
        uint32_t offset = out_v.size();
        uint32_t ioffset = out_i.size();
        // vert_count = 24;
        out_v.push_back({{1.0f, 1.0f, 1.0f}, {1.0f, 0.0f, 0.0f}, {0.0f, 0.0f}});
        out_v.push_back({{1.0f, 1.0f, -1.0f}, {1.0f, 0.0f, 0.0f}, {0.0f, 1.0f}});
        out_v.push_back({{1.0f, -1.0f, -1.0f}, {1.0f, 0.0f, 0.0f}, {1.0f, 1.0f}});
        out_v.push_back({{1.0f, -1.0f, 1.0f}, {1.0f, 0.0f, 0.0f}, {1.0f, 0.0f}});
        out_i.insert(out_i.end(), {0 + offset, 1 + offset, 2 + offset, 2 + offset, 3 + offset, 0 + offset});
        out_v.push_back({{1.0f, -1.0f, 1.0f}, {0.0f, -1.0f, 0.0f}, {0.0f, 0.0f}});
        out_v.push_back({{1.0f, -1.0f, -1.0f}, {0.0f, -1.0f, 0.0f}, {0.0f, 1.0f}});
        out_v.push_back({{-1.0f, -1.0f, -1.0f}, {0.0f, -1.0f, 0.0f}, {1.0f, 1.0f}});
        out_v.push_back({{-1.0f, -1.0f, 1.0f}, {0.0f, -1.0f, 0.0f}, {1.0f, 0.0f}});
        out_i.insert(out_i.end(), {4 + offset, 5 + offset, 6 + offset, 6 + offset, 7 + offset, 4 + offset});
        out_v.push_back({{-1.0f, -1.0f, 1.0f}, {-1.0f, 0.0f, 0.0f}, {0.0f, 0.0f}});
        out_v.push_back({{-1.0f, -1.0f, -1.0f}, {-1.0f, 0.0f, 0.0f}, {0.0f, 1.0f}});
        out_v.push_back({{-1.0f, 1.0f, -1.0f}, {-1.0f, 0.0f, 0.0f}, {1.0f, 1.0f}});
        out_v.push_back({{-1.0f, 1.0f, 1.0f}, {-1.0f, 0.0f, 0.0f}, {1.0f, 0.0f}});
        out_i.insert(out_i.end(), {8 + offset, 9 + offset, 10 + offset, 10 + offset, 11 + offset, 8 + offset});
        out_v.push_back({{-1.0f, 1.0f, 1.0f}, {0.0f, 1.0f, 0.0f}, {0.0f, 0.0f}});
        out_v.push_back({{-1.0f, 1.0f, -1.0f}, {0.0f, 1.0f, 0.0f}, {0.0f, 1.0f}});
        out_v.push_back({{1.0f, 1.0f, -1.0f}, {0.0f, 1.0f, 0.0f}, {1.0f, 1.0f}});
        out_v.push_back({{1.0f, 1.0f, 1.0f}, {0.0f, 1.0f, 0.0f}, {1.0f, 0.0f}});
        out_i.insert(out_i.end(), {12 + offset, 13 + offset, 14 + offset, 14 + offset, 15 + offset, 12 + offset});
        out_v.push_back({{-1.0f, -1.0f, 1.0f}, {0.0f, 0.0f, 1.0f}, {0.0f, 0.0f}});
        out_v.push_back({{-1.0f, 1.0f, 1.0f}, {0.0f, 0.0f, 1.0f}, {0.0f, 1.0f}});
        out_v.push_back({{1.0f, 1.0f, 1.0f}, {0.0f, 0.0f, 1.0f}, {1.0f, 1.0f}});
        out_v.push_back({{1.0f, -1.0f, 1.0f}, {0.0f, 0.0f, 1.0f}, {1.0f, 0.0f}});
        out_i.insert(out_i.end(), {16 + offset, 17 + offset, 18 + offset, 18 + offset, 19 + offset, 16 + offset});
        out_v.push_back({{-1.0f, 1.0f, -1.0f}, {0.0f, 0.0f, -1.0f}, {0.0f, 0.0f}});
        out_v.push_back({{-1.0f, -1.0f, -1.0f}, {0.0f, 0.0f, -1.0f}, {0.0f, 1.0f}});
        out_v.push_back({{1.0f, -1.0f, -1.0f}, {0.0f, 0.0f, -1.0f}, {1.0f, 1.0f}});
        out_v.push_back({{1.0f, 1.0f, -1.0f}, {0.0f, 0.0f, -1.0f}, {1.0f, 0.0f}});
        out_i.insert(out_i.end(), {20 + offset, 21 + offset, 22 + offset, 22 + offset, 23 + offset, 20 + offset});
        return ModelUnit{offset, (uint32_t)out_v.size() - offset, ioffset, (uint32_t)out_i.size() - ioffset};
    }

    ModelUnit gen_octahedron() {
        const float nv = .5773503f; // 1.0 / sqrt(3)
        uint32_t offset = out_v.size();
        uint32_t ioffset = out_i.size();
        // vert_count = 24;
        out_v.push_back({{1.0f, 0.0f, 0.0f}, {nv, nv, nv}});
        out_v.push_back({{0.0f, 0.0f, 1.0f}, {nv, nv, nv}});
        out_v.push_back({{0.0f, 1.0f, 0.0f}, {nv, nv, nv}});
        out_i.insert(out_i.end(), {0 + offset, 1 + offset, 2 + offset});
        out_v.push_back({{0.0f, -1.0f, 0.0f}, {nv, -nv, nv}});
        out_v.push_back({{0.0f, 0.0f, 1.0f}, {nv, -nv, nv}});
        out_v.push_back({{1.0f, 0.0f, 0.0f}, {nv, -nv, nv}});
        out_i.insert(out_i.end(), {3 + offset, 4 + offset, 5 + offset});
        out_v.push_back({{-1.0f, 0.0f, 0.0f}, {-nv, -nv, nv}});
        out_v.push_back({{0.0f, 0.0f, 1.0f}, {-nv, -nv, nv}});
        out_v.push_back({{0.0f, -1.0f, 0.0f}, {-nv, -nv, nv}});
        out_i.insert(out_i.end(), {6 + offset, 7 + offset, 8 + offset});
        out_v.push_back({{0.0f, 1.0f, 0.0f}, {-nv, nv, nv}});
        out_v.push_back({{0.0f, 0.0f, 1.0f}, {-nv, nv, nv}});
        out_v.push_back({{-1.0f, 0.0f, 0.0f}, {-nv, nv, nv}});
        out_i.insert(out_i.end(), {9 + offset, 10 + offset, 11 + offset});
        out_v.push_back({{1.0f, 0.0f, 0.0f}, {nv, nv, -nv}});
        out_v.push_back({{0.0f, 1.0f, 0.0f}, {nv, nv, -nv}});
        out_v.push_back({{0.0f, 0.0f, -1.0f}, {nv, nv, -nv}});
        out_i.insert(out_i.end(), {12 + offset, 13 + offset, 14 + offset});
        out_v.push_back({{0.0f, -1.0f, 0.0f}, {nv, -nv, -nv}});
        out_v.push_back({{1.0f, 0.0f, 0.0f}, {nv, -nv, -nv}});
        out_v.push_back({{0.0f, 0.0f, -1.0f}, {nv, -nv, -nv}});
        out_i.insert(out_i.end(), {15 + offset, 16 + offset, 17 + offset});
        out_v.push_back({{-1.0f, 0.0f, 0.0f}, {-nv, -nv, -nv}});
        out_v.push_back({{0.0f, -1.0f, 0.0f}, {-nv, -nv, -nv}});
        out_v.push_back({{0.0f, 0.0f, -1.0f}, {-nv, -nv, -nv}});
        out_i.insert(out_i.end(), {18 + offset, 19 + offset, 20 + offset});
        out_v.push_back({{0.0f, 1.0f, 0.0f}, {-nv, nv, -nv}});
        out_v.push_back({{-1.0f, 0.0f, 0.0f}, {-nv, nv, -nv}});
        out_v.push_back({{0.0f, 0.0f, -1.0f}, {-nv, nv, -nv}});
        out_i.insert(out_i.end(), {21 + offset, 22 + offset, 23 + offset});
        return ModelUnit{offset, (uint32_t)out_v.size() - offset, ioffset, (uint32_t)out_i.size() - ioffset};
    }

    ModelUnit gen_sphere(float radius, uint32_t slice, uint32_t stack) {
        uint32_t offset = out_v.size();
        uint32_t ioffset = out_i.size();
        // vert_count = 2 + (stack - 1) * slice;
        // vertices
        out_v.push_back({{0.0f, 0.0f, radius}, {0.0f, 0.0f, 1.0f}, {0.5f, 0.0f}});
        out_v.push_back({{0.0f, 0.0f, -radius}, {0.0f, 0.0f, -1.0f}, {0.5f, 1.0f}});
        float za = pi / stack;
        for(uint32_t i = 1; i < stack; ++i) {
            float z = radius * std::cos(za);
            float r = radius * std::sin(za);
            float a = 0;
            for(uint32_t s = 0; s < slice; ++s) {
                float x = r * std::cos(a);
                float y = r * std::sin(a);
                tmath::fvec3 position = {x, y, z};
                out_v.push_back({position, position * (1.0f / radius), {(x + 1.0f) * 0.5f, (1.0f - z / radius) * 0.5f}});
                a += pi2 / slice;
            }
            za += pi / stack;
        }
        // draw indices
        for(uint32_t i = 2; i < slice + 1; ++i)
            out_i.insert(out_i.end(), {0, i + 1, i});
        out_i.insert(out_i.end(), {0, 2, slice + 1});
        for(uint32_t s = 0; s < stack - 2; ++s) {
            uint32_t start = 2 + slice * s;
            for(uint32_t i = start; i < start + slice - 1; ++i)
                out_i.insert(out_i.end(), {i, i + 1, i + slice, i + slice, i + 1, i + slice + 1});
            out_i.insert(out_i.end(), {start + slice - 1, start, start + slice * 2 - 1, start + slice * 2 - 1, start, start + slice});
        }
        uint32_t start = 2 + slice * (stack - 2);
        for(uint32_t i = start; i < start + slice - 1; ++i)
            out_i.insert(out_i.end(), {i, i + 1, 1});
        out_i.insert(out_i.end(), {start + slice - 1, start, 1});
        // fix offset
        for(uint32_t i = ioffset; i < out_i.size(); ++i)
            out_i[i] += offset;
        return ModelUnit{offset, (uint32_t)out_v.size() - offset, ioffset, (uint32_t)out_i.size() - ioffset};
    }

    ModelUnit gen_cylinder(float base, float top, float height, uint32_t slice, uint32_t stack) {
        uint32_t offset = out_v.size();
        uint32_t ioffset = out_i.size();
        // vert_count = 2 + (stack + 3) * (slice + 1);
        // vertices
        out_v.push_back({{0.0f, 0.0f, height}, {0.0f, 0.0f, height}, {0.5f, 0.0f}});
        out_v.push_back({{0.0f, 0.0f, 0.0f}, {0.0f, 0.0f, -1.0f}, {0.5f, 1.0f}});
        for(uint32_t i = 0; i < slice + 1; ++i) {
            float a = (float)i / slice;
            out_v.push_back({{std::cos(a * pi2) * top, std::sin(a * pi2) * top, height}, {0.0f, 0.0f, 1.0f}, {a, 0.0f}});
        }
        for(uint32_t i = 0; i < slice + 1; ++i) {
            float a = (float)i / slice;
            out_v.push_back({{std::cos(a * pi2) * base, std::sin(a * pi2) * base, 0.0f}, {0.0f, 0.0f, -1.0f}, {a, 1.0f}});
        }
        float z = height;
        tmath::fvec3 norm = tmath::fvec3{top - base, 0, height}.cross(tmath::fvec3{0.0f, -1.0f, 0.0f}).normalize();
        for(uint32_t zi = 0; zi < stack + 1; ++zi) {
            float r = top + (base - top) * zi / stack;
            for(uint32_t i = 0; i < slice + 1; ++i) {
                float a = (float)i / slice;
                float x = std::cos(a * pi2);
                float y = std::sin(a * pi2);
                out_v.push_back({{x * r, y * r, z}, {x * norm[0] - y * norm[1], y * norm[0] + x * norm[1], norm[2]}, {a, 1.0f - z / height}});
            }
            z -= height / stack;
        }
        // draw indices
        for(uint32_t i = 2; i < 2 + slice; ++i)
            out_i.insert(out_i.end(), {0, i + 1, i});
        for(uint32_t i = 2 + slice + 1; i < 2 + slice * 2 + 1; ++i)
            out_i.insert(out_i.end(), {1, i, i + 1});
        for(uint32_t zi = 0; zi < stack; ++zi) {
            uint32_t st = 2 + (zi + 2) * (slice + 1);
            for(uint32_t i = st; i < st + slice; ++i)
                out_i.insert(out_i.end(), {i, i + 1, i + slice + 1, i + slice + 1, i + 1, i + slice + 2});
        }
        // fix offset
        for(uint32_t i = ioffset; i < out_i.size(); ++i)
            out_i[i] += offset;
        return ModelUnit{offset, (uint32_t)out_v.size() - offset, ioffset, (uint32_t)out_i.size() - ioffset};
    }

    // radius = f(z)
    template<typename F>
    ModelUnit gen_solid_of_revolution(float zb, float zt, const F& f, uint32_t slice, uint32_t stack) {
        uint32_t offset = out_v.size();
        uint32_t ioffset = out_i.size();
        // vert_count = 2 + (stack + 3) * (slice + 1);
        // vertices
        out_v.push_back({{0.0f, 0.0f, zt}, {0.0f, 0.0f, zt}, {0.5f, 0.0f}});
        out_v.push_back({{0.0f, 0.0f, zb}, {0.0f, 0.0f, zb}, {0.5f, 1.0f}});
        float top = f(zt);
        for(uint32_t i = 0; i < slice + 1; ++i) {
            float a = (float)i / slice;
            out_v.push_back({{std::cos(a * pi2) * top, std::sin(a * pi2) * top, zt}, {0.0f, 0.0f, 1.0f}, {a, 0.0f}});
        }
        float base = f(zb);
        for(uint32_t i = 0; i < slice + 1; ++i) {
            float a = (float)i / slice;
            out_v.push_back({{std::cos(a * pi2) * base, std::sin(a * pi2) * base, 0.0f}, {0.0f, 0.0f, -1.0f}, {a, 1.0f}});
        }
        rs::vector<float> rsample(stack + 1);
        rs::vector<tmath::fvec3> norms(stack + 1);
        float deltaz = (zb - zt) / stack;
        for(uint32_t i = 0; i < stack + 1; ++i)
            rsample[i] = f(zt + deltaz * i);
        norms[0] = tmath::fvec3{rsample[0] - rsample[1], 0, -deltaz}.cross(tmath::fvec3{0.0f, -1.0f, 0.0f}).normalize();
        for(uint32_t i = 1; i < stack; ++i)
            norms[i] = tmath::fvec3{rsample[i - 1] - rsample[i + 1], 0, -deltaz * 2.0f}.cross(tmath::fvec3{0.0f, -1.0f, 0.0f}).normalize();
        norms[stack] = tmath::fvec3{rsample[stack - 1] - rsample[stack], 0, -deltaz}.cross(tmath::fvec3{0.0f, -1.0f, 0.0f}).normalize();
        // tmath::fvec3 norm = tmath::fvec3{top - base, 0, height}.cross(tmath::fvec3{0.0f, -1.0f, 0.0f}).normalize();
        float z = zt;
        for(uint32_t zi = 0; zi < stack + 1; ++zi) {
            float r = rsample[zi];
            tmath::fvec3& norm = norms[zi];
            for(uint32_t i = 0; i < slice + 1; ++i) {
                float a = (float)i / slice;
                float x = std::cos(a * pi2);
                float y = std::sin(a * pi2);
                out_v.push_back({{x * r, y * r, z}, {x * norm[0] - y * norm[1], y * norm[0] + x * norm[1], norm[2]}, {a, 1.0f - deltaz * zi / (zb - zt)}});
            }
            z += deltaz;
        }
        // draw indices
        for(uint32_t i = 2; i < 2 + slice; ++i)
            out_i.insert(out_i.end(), {0, i + 1, i});
        for(uint32_t i = 2 + slice + 1; i < 2 + slice * 2 + 1; ++i)
            out_i.insert(out_i.end(), {1, i, i + 1});
        for(uint32_t zi = 0; zi < stack; ++zi) {
            uint32_t st = 2 + (zi + 2) * (slice + 1);
            for(uint32_t i = st; i < st + slice; ++i)
                out_i.insert(out_i.end(), {i, i + 1, i + slice + 1, i + slice + 1, i + 1, i + slice + 2});
        }
        // fix offset
        for(uint32_t i = ioffset; i < out_i.size(); ++i)
            out_i[i] += offset;
        return ModelUnit{offset, (uint32_t)out_v.size() - offset, ioffset, (uint32_t)out_i.size() - ioffset};
    }

    uint32_t allocate_attribute_index() {
        if(empty_indices.empty()) {
            attrs.emplace_back();
            return max_index++;
        } else {
            auto index = empty_indices.back();
            empty_indices.pop_back();
            return index;
        }
    }

    ModelUnit clone_unit(ModelUnit pu) {
        uint32_t new_voffset = out_v.size();
        uint32_t new_ioffset = out_i.size();
        out_v.insert(out_v.end(), &out_v[pu.vert_offset], &out_v[pu.vert_offset + pu.vert_count]);
        out_i.insert(out_i.end(), &out_i[pu.index_offset], &out_i[pu.index_offset + pu.index_count]);
        uint32_t diff = new_ioffset - pu.index_offset;
        for(uint32_t i = new_ioffset; i < out_i.size(); ++i)
            out_i[i] -= diff;
    }

    rs::vector<Vertex>& all_vertices() { return out_v; }
    rs::vector<uint32_t>& all_indices() { return out_i; }
    rs::vector<VertexAttribute>& all_attributes() { return attrs; }

private:
    rs::vector<Vertex> out_v;
    rs::vector<uint32_t> out_i;
    rs::vector<uint32_t> empty_indices;
    uint32_t max_index = 0;
    rs::vector<VertexAttribute> attrs;
};

#endif