#pragma once

#include <optional>
#include "./core.hpp"

namespace generational {

    /// Generational container accessible by id
    template<typename T, typename S = 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);
            }
        }

        key<S> push(const T& in) {
            if(freed.empty()) {
                const auto idx = entries.size();
                entries.emplace_back(in);
                return id(idx, 0);
            } else {
                const auto idx = freed.back();
                freed.pop_back();
                auto &entry = entries[idx];
                assert(!entry.value.has_value());
                entry.value = in;
                return entry.next(idx, 0);
            }
        }
        template <typename... _Args>
        std::pair<key<S>, T*> emplace(uint8_t flag, _Args &&... __args) {
            if(freed.empty()) {
                const auto idx = entries.size();
                auto& entry = entries.emplace_back(entry::merge(0, flag), std::forward<_Args>(__args)...);
                return std::make_pair(id(idx, 0, flag), &entry.value.value());
            } 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 std::make_pair(entry.next(idx, flag), &entry.value.value());
            }
        }
        template <typename... _Args>
        T& set_emplace(key<S> at, _Args &&... __args) {
            const auto idx = at.val.index();
            if (idx >= entries.size())
                entries.resize(idx+1);

            auto& entry = entries.at(idx);
            assert(!entry.value.has_value());
            entry.val = entry::merge(at.val.generation(), at.val.flag());
            entry.value.emplace(std::forward<_Args>(__args)...);
            return entry.value.value();
        }
        bool put(key<S> k, const T& in) {
            const auto idx = k.val.index();
            if (idx >= entries.size())
                return false;

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

            entries[idx].value = in;
            return true;
        }
        T& at(key<S> k) {
            assert(contains(k));
            return entries[k.val.index()].value.value();
        }
        T* directly_at(key<S> k) {
            if (k.val.index() >= entries.size())
                return nullptr;

            auto &cell = entries.at(k.val.index());
            if (!cell.value.has_value() || cell.generation() != k.val.generation())
                return nullptr;

            return &cell.value.value();
        }
        const T* directly_at(key<S> k) const {
            const auto idx = k.val.index();
            if (idx >= entries.size())
                return nullptr;

            auto &cell = entries.at(idx);
            if (!cell.value.has_value() || cell.generation() != k.val.generation())
                return nullptr;

            return &cell.value.value();
        }
        T* directly_at(ref<S> index) {
            if (index.val >= entries.size() || !entries.at(index.val).value.has_value())
                return nullptr;

            return &entries[index.val].value.value();
        }
        const T* directly_at(ref<S> index) const {
            if (index.val >= entries.size() || !entries.at(index.val).value.has_value())
                return nullptr;

            return &entries[index.val].value.value();
        }

        bool contains(key<S> k) const {
            const auto idx = k.val.index();
            return idx < entries.size() &&
                   entries.at(idx).generation() == k.val.generation() &&
                   entries.at(idx).value.has_value();
        }

        /// Vanished idx will not be considered free (and so not reused)
        bool vanish(key<S> k) {
            if(!contains(k))
                return false;

            const auto idx = k.val.index();
            entries[idx].value = std::nullopt;
            return true;
        }
        bool free(key<S> k) {
            if(!contains(k))
                return false;

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

        bool empty() const {
            for(auto& v: entries) {
                if(v.value.has_value())
                    return false;
            }
            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(key<S>(id(i, entry.generation(), entry.flag())), entry.value.value());
                }
            }
        }

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

        template<typename predicate>
        bool any(predicate fn) const {
            for (size_t i = 0; i < entries.size(); i++) {
                const auto &entry = entries[i];
                if(entry.value.has_value() && fn(key<S>(id(i, entry.generation(), entry.flag())), entry.value.value())) {
                    return true;
                }
            }
            return false;
        }

        /// See vanish
        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(key<S>(id(i, entry.generation(), entry.flag())), 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(key<S>(id(i, entry.generation(), entry.flag())), 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(); });
        }
        void reserve(size_t n) { entries.reserve(n); }

        key<S> with_flag(const ref<S>& r) const {
            if (r.val >= entries.size())
                return id();

            const auto &cell = entries.at(r.val);
            if (!cell.value.has_value())
                return id();

            return id(r.val, cell.generation(), cell.flag());
        }
        key<S> with_flag(const key<S>& k) const {
            const auto idx = k.val.index();
            if (idx >= entries.size())
                return id();

            const auto &cell = entries.at(idx);
            const auto gen = k.val.generation();
            if (!cell.value.has_value() || cell.generation() != gen)
                return id();

            return id(idx, gen, cell.flag());
        }
        void set_flag(const key<S>& k) {
            const auto idx = k.val.index();
            if (idx >= entries.size())
                return;

            entries.at(idx).set_flag(k.val.flag());
        }

    private:
        struct entry {
            entry(): value(std::nullopt), val(0) { }
            entry(const T &in, uint_fast32_t val = 0):
                value(in), val(val) { }
            template <typename... _Args>
            entry(int_fast32_t val, _Args &&... __args): val(val) {
                value.emplace(std::forward<_Args>(__args)...);
            }
            std::optional<T> value;
            uint_fast32_t val = 0;

            static constexpr uint32_t merge(uint_fast32_t gen, uint_fast8_t flag) {
                return (gen & id::GENERATION_MASK) | ((flag & id::FLAG_MASK) << FLAG_OFFSET);
            }
            constexpr id::type generation() const { return val & id::GENERATION_MASK; }
            static constexpr size_t FLAG_OFFSET = id::FLAG_OFFSET - id::GENERATION_OFFSET;
            constexpr uint_fast8_t flag() const { return (val >> FLAG_OFFSET) & id::FLAG_MASK; }
            inline constexpr void set_flag(uint_fast8_t flag) { val = merge(val & ~id::FLAG_MASK, flag); }
            inline key<S> next(size_t idx, uint8_t flag) {
                val = merge(++val & ~id::FLAG_MASK, flag);
                return id(idx, generation(), flag);
            }
        };
        std::vector<entry> entries;
        std::vector<uint_fast32_t> freed;
    };

    /// Dummy generational container without allocation nor validation. A view to vector
    template<typename F, typename T = F>
    class view_vector: public std::vector<std::optional<T>> {
    public:
        view_vector(): std::vector<std::optional<T>>() { }

        template <typename... _Args>
        void set_emplace(ref<F> r, _Args &&... __args) {
            if(r.val >= this->size())
                this->resize(r.val + 1);

            this->at(r.val).emplace(std::forward<_Args>(__args)...);
        }
        void set(ref<F> r, const T& in) {
            if(r.val >= this->size())
                this->resize(r.val + 1);

            this->at(r.val) = in;
        }

        T* get(ref<F> r) {
            if (r.val >= this->size() || !this->at(r.val).has_value())
                return nullptr;

            return &this->at(r.val).value();
        }
        const T* get(ref<F> r) const {
            if (r.val >= this->size() || !this->at(r.val).has_value())
                return nullptr;

            return &this->at(r.val).value();
        }

        template<typename apply>
        void iter(apply fn) const {
            for (size_t i = 0; i < this->size(); i++) {
                if(this->at(i).has_value()) {
                    fn(ref<F>(i), this->at(i).value());
                }
            }
        }

        template<typename predicate>
        bool contains(predicate fn) const {
            for (size_t i = 0; i < this->size(); i++) {
                if(this->at(i).has_value() && fn(ref<F>(i), this->at(i).value())) {
                    return true;
                }
            }
            return false;
        }

        void erase(ref<F> r) {
            if(r.val < this->size())
                this->at(r.val) = std::nullopt;
        }
        template<typename extractor>
        void remove(extractor fn) {
            for (size_t i = 0; i < this->size(); i++) {
                if(this->at(i).has_value() && fn(ref<F>(i), this->at(i).value())) {
                    erase(i);
                }
            }
        }

        size_t count() const {
            return std::count_if(this->begin(), this->end(),
                [](const std::optional<T> &e) { return e.has_value(); });
        }
    };

}