#pragma once
#include <algorithm>
#include <concepts>
#include <cstdint>
#include <cstdlib>
#include <format>
#include <set>
#include <stdexcept>
#include <type_traits>
#include <unordered_map>
#include <vector>

namespace zecsy
{
    using entity_id = uint64_t;

    template<typename... T>
    concept Component = []
    {
        static_assert((std::is_default_constructible_v<T> && ...),
                      "Should have a default constructor");
        static_assert((std::is_trivially_copyable_v<T> && ...),
                      "Should be trivially copyable");
        static_assert((std::is_trivially_destructible_v<T> && ...),
                      "Should be trivially destructible");
        static_assert((std::is_standard_layout_v<T> && ...),
                      "Should have standard layout");
        return true;
    }();

    class world final
    {
    public:
        entity_id make_entity();
        void destroy_entity(entity_id e);
        bool is_alive(entity_id e) const;
        size_t components_in_entity(entity_id e) const;

        size_t entity_count() const;
        size_t total_component_count() const;
        size_t archetype_count() const;

        template<Component T>
        size_t component_count();

        template<Component T>
        bool has(entity_id e) const;

        template<Component First, Component Second, Component... Rest>
        bool has(entity_id e) const;

        template<Component T>
        T& get(entity_id e);

        template<Component T>
        void set(entity_id e);

        template<Component T>
        void set(entity_id e, const T& comp);

        template<Component First, Component Second, Component... Rest>
        void set(entity_id e);

        template<Component First, Component Second, Component... Rest>
        void set(entity_id e, const First& comp0, const Second& comp1,
                 const Rest&... rest_comps);

        template<Component T>
        void remove(entity_id e);

        template<Component First, Component Second, Component... Rest>
        void remove(entity_id e);

        template<Component... T>
        void query(std::invocable<entity_id, T&...> auto&& system);

        size_t get_archetypes_checked() const;
        size_t get_entities_processed() const;

    private:
        using comp_id = size_t;
        std::set<entity_id> alive_entities;
        std::unordered_map<entity_id, std::set<comp_id>> entity_to_comps;
        entity_id entity_counter = 0;

        size_t query_archetypes_checked = 0;
        size_t query_entities_processed = 0;

        struct component_pool
        {
            std::vector<uint8_t> data;
            std::vector<size_t> free_list;
            std::unordered_map<entity_id, size_t> entity_to_index;
            std::unordered_map<size_t, entity_id> index_to_entity;
        };

        std::unordered_map<comp_id, component_pool> pools;

        using archetype_signature = std::vector<comp_id>;
        using entity_group = std::set<entity_id>;

        struct archetype_hash
        {
            size_t operator()(const archetype_signature& vec) const;
        };

        std::unordered_map<archetype_signature, entity_group, archetype_hash>
            archetypes;

        template<Component T>
        static comp_id get_component_id();

        static comp_id next_component_id;
    };

    inline size_t world::archetype_hash::operator()(
        const world::archetype_signature& vec) const
    {
        size_t seed = vec.size();
        for(const auto& id: vec)
        {
            seed ^= id + 0x9e3779b9 + (seed << 6) + (seed >> 2);
        }
        return seed;
    }

    template<Component T>
    inline world::comp_id world::get_component_id()
    {
        static comp_id id = next_component_id++;
        return id;
    }

    inline world::comp_id world::next_component_id = 0;

    inline size_t world::components_in_entity(entity_id e) const
    {
        return entity_to_comps.contains(e) ? entity_to_comps.at(e).size() : 0;
    }

    inline size_t world::entity_count() const
    {
        return alive_entities.size();
    }

    inline size_t world::total_component_count() const
    {
        size_t count = 0;

        for(const auto& [id, pool]: pools)
        {
            count += pool.entity_to_index.size();
        }

        return count;
    }

    inline size_t world::archetype_count() const
    {
        return archetypes.size();
    }

    template<Component T>
    inline size_t world::component_count()
    {
        const comp_id id = get_component_id<T>();
        const auto it = pools.find(id);
        return it != pools.end() ? it->second.entity_to_index.size() : 0;
    }

    inline size_t world::get_archetypes_checked() const
    {
        return query_archetypes_checked;
    }

    inline size_t world::get_entities_processed() const
    {
        return query_entities_processed;
    }

    inline entity_id world::make_entity()
    {
        auto id = ++entity_counter;
        alive_entities.emplace(id);
        entity_to_comps[id] = {};

        std::vector<comp_id> key;
        auto& group = archetypes[key];
        group.emplace(id);

        return id;
    }

    inline void world::destroy_entity(entity_id e)
    {
        alive_entities.erase(e);

        auto& comp_set = entity_to_comps[e];
        std::vector<comp_id> key(comp_set.begin(), comp_set.end());

        auto& group = archetypes[key];
        group.erase(e);

        if(archetypes[key].empty())
        {
            archetypes.erase(key);
        }

        for(comp_id id: comp_set)
        {
            auto& pool = pools[id];
            auto index = pool.entity_to_index[e];
            pool.entity_to_index.erase(e);
            pool.index_to_entity.erase(index);
            pool.free_list.emplace_back(index);
        }

        entity_to_comps.erase(e);
    }

    inline bool world::is_alive(entity_id e) const
    {
        return alive_entities.contains(e);
    }

    template<Component T>
    inline bool world::has(entity_id e) const
    {
        if(entity_to_comps.contains(e))
        {
            return entity_to_comps.at(e).contains(get_component_id<T>());
        }
        return false;
    }

    template<Component T>
    inline T& world::get(entity_id e)
    {
        auto id = get_component_id<T>();
        if(!has<T>(e))
        {
            throw std::runtime_error(
                std::format("Entity #{} doesn't have {}", e, typeid(T).name()));
        }

        auto& pool = pools.at(id);
        auto index = pool.entity_to_index.at(e);
        return *reinterpret_cast<T*>(&pool.data[index * sizeof(T)]);
    }

    template<Component T>
    inline void world::set(entity_id e, const T& comp)
    {
        auto id = get_component_id<T>();
        auto& pool = pools[id];

        size_t index;
        if(!pool.free_list.empty())
        {
            index = pool.free_list.back();
            pool.free_list.pop_back();
        }
        else
        {
            index = pool.data.size() / sizeof(T);
            pool.data.resize(pool.data.size() + sizeof(T));
        }

        new(&pool.data[index * sizeof(T)]) T(comp);
        pool.entity_to_index[e] = index;
        pool.index_to_entity[index] = e;

        auto& comp_set = entity_to_comps[e];
        std::set<comp_id> old_set = comp_set;
        auto [it, inserted] = comp_set.insert(id);

        if(inserted)
        {
            std::vector<comp_id> old_key(old_set.begin(), old_set.end());

            auto& group = archetypes[old_key];
            group.erase(e);

            if(archetypes[old_key].empty())
            {
                archetypes.erase(old_key);
            }

            std::vector<comp_id> new_key(comp_set.begin(), comp_set.end());
            archetypes[new_key].emplace(e);
        }
    }

    template<Component T>
    inline void world::set(entity_id e)
    {
        set(e, T{});
    }

    template<Component T>
    inline void world::remove(entity_id e)
    {
        auto id = get_component_id<T>();
        auto& comp_set = entity_to_comps[e];

        if(comp_set.erase(id) > 0)
        {
            std::vector<comp_id> old_key(comp_set.begin(), comp_set.end());
            old_key.push_back(id);
            std::sort(old_key.begin(), old_key.end());

            auto& old_group = archetypes[old_key];
            old_group.erase(e);

            std::vector<comp_id> new_key(comp_set.begin(), comp_set.end());
            std::sort(new_key.begin(), new_key.end());

            archetypes[new_key].emplace(e);

            if(archetypes[old_key].empty())
            {
                archetypes.erase(old_key);
            }

            auto& pool = pools[id];
            auto index = pool.entity_to_index[e];
            pool.free_list.push_back(index);
            pool.entity_to_index.erase(e);
            pool.index_to_entity.erase(index);
        }
    }

    template<Component First, Component Second, Component... Rest>
    inline bool world::has(entity_id e) const
    {
        return has<First>(e) && has<Second>(e) && (has<Rest>(e) && ...);
    }

    template<Component First, Component Second, Component... Rest>
    inline void world::set(entity_id e)
    {
        set(e, First{});
        set(e, Second{});
        (set(e, Rest{}), ...);
    }

    template<Component First, Component Second, Component... Rest>
    inline void world::set(entity_id e, const First& comp0, const Second& comp1,
                           const Rest&... rest_comps)
    {
        set(e, comp0);
        set(e, comp1);
        (set(e, rest_comps), ...);
    }

    template<Component First, Component Second, Component... Rest>
    inline void world::remove(entity_id e)
    {
        remove<First>(e);
        remove<Second>(e);
        (remove<Rest>(e), ...);
    }

    template<Component... T>
    inline void world::query(std::invocable<entity_id, T&...> auto&& system)
    {
        std::vector<comp_id> required = {get_component_id<T>()...};
        std::sort(required.begin(), required.end());

        query_archetypes_checked = 0;
        query_entities_processed = 0;

        for(const auto& [archetype_key, entities]: archetypes)
        {
            query_archetypes_checked++;

            bool match = true;
            for(comp_id req_id: required)
            {
                if(!std::binary_search(archetype_key.begin(),
                                       archetype_key.end(), req_id))
                {
                    match = false;
                    break;
                }
            }

            if(match)
            {
                query_entities_processed += entities.size();
                for(entity_id e: entities)
                {
                    system(e, get<T>(e)...);
                }
            }
        }
    }
}; // namespace zecsy