#include <iostream>
#include <functional>
#include <typeindex>
#include <typeinfo>
#include <unordered_map>
#include <fmt/core.h>
#include <any>

using namespace fmt;

struct Position {
  double x, y;
};

struct Velocity {
  double x, y;
};

typedef unsigned long Entity;

typedef std::unordered_map<Entity, std::any> EntityMap;

struct World {
  unsigned long entity_count = 0;
  std::unordered_map<std::type_index, EntityMap> $components;

  Entity entity() {
    return ++entity_count;
  }

  template <typename Comp>
  EntityMap& entity_map_for() {
    return $components[std::type_index(typeid(Comp))];
  }

  template <typename Comp>
  void assign(Entity ent, Comp val) {
    EntityMap &map = entity_map_for<Comp>();
    map[ent] = val;
  }

  template <typename Comp>
  Comp &component(Entity ent) {
    EntityMap &map = entity_map_for<Comp>();
    std::any &res = map[ent];
    return std::any_cast<Comp&>(res);
  }

  template<typename Comp>
  void system(std::function<void(const Entity&, Comp&)> cb) {
    EntityMap &map = entity_map_for<Comp>();
    for(auto& [entity, any_comp] : map) {
      Comp &res = std::any_cast<Comp&>(any_comp);
      cb(entity, res);
    }
  }

  template<typename Comp>
  void remove(Entity ent) {
    EntityMap &map = entity_map_for<Comp>();
    map.erase(ent);
  }

  template<typename CompA, typename CompB>
  void system(std::function<void(const Entity&, CompA&, CompB&)> cb) {
    EntityMap &map_a = entity_map_for<CompA>();
    EntityMap &map_b = entity_map_for<CompB>();

    for(auto& [entity, any_a] : map_a) {
      if(map_b.contains(entity)) {
        CompA &res_a = std::any_cast<CompA&>(any_a);
        CompB &res_b = component<CompB>(entity);
        cb(entity, res_a, res_b);
      }
    }
  }
};

int main() {
  World me;

  Entity test = me.entity();
  Entity test2 = me.entity();

  me.assign<Position>(test, {10,20});
  me.assign<Velocity>(test, {1,2});

  me.assign<Position>(test2, {1,1});
  me.assign<Velocity>(test2, {10,20});

  Position &pos = me.component<Position>(test);
  println("GOT POS x={}, y={}", pos.x, pos.y);

  Velocity &vel = me.component<Velocity>(test);
  println("GOT VELOCITY x={}, y={}", vel.x, vel.y);

  println("--- Position only system:");
  me.system<Position>([](const auto &ent, auto &pos) {
      println("entity={}, pos.x={}, pos.y={}", ent, pos.x, pos.y);
  });

  println("--- Velocity only system:");
  me.system<Velocity>([](const auto &, auto &vel) {
      println("vel.x={}, vel.y={}", vel.x, vel.y);
  });

  println("--- Manually get the velocity in position system:");
  me.system<Position>([&](const auto &ent, auto &pos) {
      Velocity &vel = me.component<Velocity>(ent);
      println("entity={}, vel.x, vel.y, pos.x={}, pos.y={}", ent, vel.x, vel.y, pos.x, pos.y);
  });

  println("--- Query only entities with Position and Velocity:");
  me.system<Position, Velocity>([&](const auto &ent, auto &pos, auto &vel) {
      println("entity={}, vel.x, vel.y, pos.x={}, pos.y={}", ent, vel.x, vel.y, pos.x, pos.y);
  });

  // now remove Velocity
  me.remove<Velocity>(test);

  println("--- After remove test, should only result in test2:");
  me.system<Position, Velocity>([&](const auto &ent, auto &pos, auto &vel) {
      println("entity={}, vel.x, vel.y, pos.x={}, pos.y={}", ent, vel.x, vel.y, pos.x, pos.y);
  });

  return 0;
}