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

#ifndef ZECSY_MAX_COMPONENTS
#define ZECSY_MAX_COMPONENTS 32
#endif // !ZECSY_MAX_COMPONENTS

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>
        T& ensure(entity_id e);

        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;
        using entity_group = std::set<entity_id>;
        using archetype_signature = std::bitset<ZECSY_MAX_COMPONENTS>;

        std::unordered_map<entity_id, archetype_signature> 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<comp_id, component_pool> pools;
        std::unordered_map<archetype_signature, entity_group> archetypes;

        template<Component T>
        static comp_id get_component_id();

        static comp_id next_component_id;
    };

    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).count() : 0;
    }

    inline size_t world::entity_count() const
    {
        return entity_to_comps.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;
        entity_to_comps[id] = {};

        archetype_signature key;
        auto& group = archetypes[key];
        group.emplace(id);

        return id;
    }

    inline void world::destroy_entity(entity_id e)
    {
        auto archetype = entity_to_comps[e];

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

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

        for(int id = 0; id < ZECSY_MAX_COMPONENTS; ++id)
        {
            if(archetype.test(id))
            {
                auto& pool = pools[id];
                auto index = pool.entity_to_index[e];
                pool.entity_to_index.erase(e);
                pool.free_list.emplace_back(index);
            }
        }

        entity_to_comps.erase(e);
    }

    inline bool world::is_alive(entity_id e) const
    {
        return entity_to_comps.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).test(get_component_id<T>());
        }
        return false;
    }

    template<Component T>
    T& world::ensure(entity_id e)
    {
        if(!has<T>(e))
        {
            set<T>(e);
        }

        return get<T>(e);
    }

    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;

        auto& archetype = entity_to_comps[e];
        auto old_archetype = archetype;
        archetype.set(id);

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

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

        archetypes[archetype].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)
    {
        if(!has<T>(e))
        {
            return;
        }

        auto id = get_component_id<T>();
        auto& archetype = entity_to_comps[e];

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

        if(old_group.empty())
        {
            archetypes.erase(archetype);
        }

        archetype.reset(id);
        archetypes[archetype].emplace(e);

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

    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>()...};*/

        archetype_signature required;

        (required.set(get_component_id<T>()), ...);

        query_archetypes_checked = 0;
        query_entities_processed = 0;

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

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