#pragma once

#include <cassert>
#include <vector>
#include <optional>

namespace data::generational {
    struct id {
        id(size_t index, size_t generation = 0): index(index), generation(generation) { }
        id(): id(0) { }
        size_t index;
        size_t generation;

        bool operator==(const id &i) const { return index == i.index && generation == i.generation; }
    };

    class allocator {
    public:
        allocator() { }
        template <typename C>
        allocator(const C& map) {
            for(const auto& [key, _]: map) {
                if (const auto size = entries.size(); key >= size) {
                    entries.resize(key + 1);
                    for (size_t i = size; i < entries.size(); i++) {
                        entries[i].is_live = false;
                        freed.push_back(i);
                    }
                }
                assert(!entries[key].is_live);
                entries[key].is_live = true;
            }
        }
        id alloc() {
            if(freed.empty()) {
                const auto idx = entries.size();
                entries.emplace_back();
                return id(idx);
            } else {
                const auto idx = freed.back();
                freed.pop_back();
                auto &entry = entries[idx];
                assert(!entry.is_live);
                entry.is_live = true;
                return id(idx, ++entries[idx].generation);
            }
        }
        bool is_live(id id) const {
            return id.index < entries.size() &&
                entries[id.index].is_live &&
                entries[id.index].generation == id.generation;
        }
        bool free(id id) {
            if(!is_live(id))
                return false;

            entries[id.index].is_live = false;
            freed.push_back(id.index);
            return true;
        }

    private:
        struct entry {
            bool is_live = true;
            size_t generation = 0;
        };
        std::vector<entry> entries;
        std::vector<size_t> freed;
    };

    template<typename T>
    class vector {
    public:
        vector() { }
        template<typename C>
        vector(const C& map) {
            for(const auto& [key, value]: map) {
                if (const auto size = entries.size(); key >= size) {
                    entries.resize(key + 1);
                }
                entries[key].value = value;
            }
            for (size_t i = 0; i < entries.size(); i++) {
                if(!entries[i].value.has_value())
                    freed.push_back(i);
            }
        }

        id push(const T& in) {
            if(freed.empty()) {
                const auto idx = entries.size();
                entries.emplace_back(in);
                return id(idx);
            } else {
                const auto idx = freed.back();
                freed.pop_back();
                auto &entry = entries[idx];
                assert(!entry.value.has_value());
                entry.value = in;
                return id(idx, ++entries[idx].generation);
            }
        }
        template <typename... _Args>
        id emplace(_Args &&... __args) {
            if(freed.empty()) {
                const auto idx = entries.size();
                entries.emplace_back(std::forward<_Args>(__args)...);
                return id(idx);
            } else {
                const auto idx = freed.back();
                freed.pop_back();
                auto &entry = entries[idx];
                assert(!entry.value.has_value());
                entry.value.emplace(std::forward<_Args>(__args)...);
                return id(idx, ++entries[idx].generation);
            }
        }
        bool put(id idx, const T& in) {
            if(idx.index >= entries.size())
                return false;

            if(entries[idx.index].generation != idx.generation || entries[idx.index].value.has_value())
                return false;

            entries[idx.index].value = in;
            return true;
        }
        T& at(id idx) {
            assert(contains(idx));
            return entries[idx.index].value.value();
        }

        bool contains(id idx) const {
            return idx.index < entries.size() &&
                   entries[idx.index].generation == idx.generation &&
                   entries[idx.index].value.has_value();
        }

        bool free(id idx) {
            if(!contains(idx))
                return false;

            entries[idx.index].value = std::nullopt;
            freed.push_back(idx.index);
            return true;
        }

        template<typename apply>
        void iter(apply fn) const {
            for (size_t i = 0; i < entries.size(); i++) {
                const auto &entry = entries[i];
                if(entry.value.has_value()) {
                    fn(id(i, entry.generation), entry.value.value());
                }
            }
        }

        template<typename apply>
        void for_each(apply fn) {
            for (size_t i = 0; i < entries.size(); i++) {
                auto &entry = entries[i];
                if(entry.value.has_value()) {
                    fn(id(i, entry.generation), entry.value.value());
                }
            }
        }

        template<typename extractor>
        void extract(extractor fn) {
            for (size_t i = 0; i < entries.size(); i++) {
                auto &entry = entries[i];
                if(entry.value.has_value()) {
                    if(fn(id(i, entry.generation), entry.value.value()))
                        entry.value = std::nullopt;
                }
            }
        }

        template<typename extractor>
        void remove(extractor fn) {
            for (size_t i = 0; i < entries.size(); i++) {
                auto &entry = entries[i];
                if(entry.value.has_value()) {
                    if(fn(id(i, entry.generation), entry.value.value())) {
                        entry.value = std::nullopt;
                        freed.push_back(i);
                    }
                }
            }
        }

        size_t size() const {
            return std::count_if(entries.begin(), entries.end(),
                [](const entry &e) { return e.value.has_value(); });
        }

    private:
        struct entry {
            entry(): value(std::nullopt), generation(0) { }
            entry(const T &in, size_t gen = 0):
                value(in), generation(gen) { }
            template <typename... _Args>
            entry(_Args &&... __args): generation(0) {
                value.emplace(std::forward<_Args>(__args)...);
            }
            std::optional<T> value;
            size_t generation;
        };
        std::vector<entry> entries;
        std::vector<size_t> freed;
    };
}

namespace std {
    template<>
    struct hash<data::generational::id> {
        std::size_t operator()(const data::generational::id& i) const noexcept {
            return std::hash<size_t>{}(i.index);
        }
    };
}