#ifndef _SHARED_MAP_H_
#define _SHARED_MAP_H_

#include "shared/avl_tree.h"
#include "shared/stdmacro.h"

#define MAP_PASTE_KV(K, V) CONCAT(CONCAT(K, ___), V)
#define map_entry_t(K, V) CONCAT(map_entry_t__, MAP_PASTE_KV(K, V))
#define map_t(K, V) CONCAT(map_t__, MAP_PASTE_KV(K, V))

#define map_new(K, V) CONCAT(map_new__, MAP_PASTE_KV(K, V))
#define map_free(K, V) CONCAT(map_free__, MAP_PASTE_KV(K, V))
#define map_put(K, V) CONCAT(map_put__, MAP_PASTE_KV(K, V))
#define map_get(K, V) CONCAT(map_get__, MAP_PASTE_KV(K, V))
#define map_erase(K, V) CONCAT(map_erase__, MAP_PASTE_KV(K, V))
#define _map_compare_(K, V) CONCAT(_map_compare__, MAP_PASTE_KV(K, V))
#define _map_dtor_(K, V) CONCAT(_map_dtor__, MAP_PASTE_KV(K, V))

#define MAP_DECL(K, V)                                 \
  typedef struct {                                     \
    K key;                                             \
    V value;                                           \
  } map_entry_t(K, V);                                 \
  AVL_TREE_DECL(map_entry_t(K, V));                    \
  typedef struct {                                     \
    avl_tree_t(map_entry_t(K, V)) * tree;              \
  } map_t(K, V);                                       \
  map_t(K, V) * map_new(K, V)();                       \
  void map_free(K, V)(map_t(K, V) * map);              \
  V* map_get(K, V)(map_t(K, V) * map, K key);          \
  V* map_put(K, V)(map_t(K, V) * map, K key, V value); \
  bool map_erase(K, V)(map_t(K, V) * map, K key, V * out);

#define MAP_IMPL(K, V, CMP, DTOR)                                              \
  static inline int _map_compare_(K, V)(                                       \
      map_entry_t(K, V) e1, map_entry_t(K, V) e2)                              \
  {                                                                            \
    return CMP(e1.key, e2.key);                                                \
  }                                                                            \
  static inline void _map_dtor_(K, V)(map_entry_t(K, V) e1)                    \
  {                                                                            \
    DTOR(e1.value);                                                            \
  }                                                                            \
  AVL_TREE_IMPL(map_entry_t(K, V), _map_compare_(K, V), _map_dtor_(K, V))      \
  map_t(K, V) * map_new(K, V)()                                                \
  {                                                                            \
    map_t(K, V)* ret = ALLOC(sizeof(map_t(K, V)));                             \
    ret->tree = avl_tree_new(map_entry_t(K, V))();                             \
    return ret;                                                                \
  }                                                                            \
  void map_free(K, V)(map_t(K, V) * map) { \
    avl_tree_free(map_entry_t(K, V))(map->tree); \
    FREE(map); \
  }\
  V* map_get(K, V)(map_t(K, V) * map, K key)                                   \
  {                                                                            \
    map_entry_t(K, V) key_;                                                    \
    key_.key = key;                                                            \
    map_entry_t(K, V)* entry =                                                 \
        avl_tree_find(map_entry_t(K, V))(map->tree, key_);                     \
                                                                               \
    if (!entry) {                                                              \
      return NULL;                                                             \
    }                                                                          \
                                                                               \
    return &entry->value;                                                      \
  }                                                                            \
  V* map_put(K, V)(map_t(K, V) * map, K key, V val)                            \
  {                                                                            \
    map_entry_t(K, V) entry;                                                   \
    entry.key = key;                                                           \
    entry.value = val;                                                         \
    map_entry_t(K, V)* ref =                                                   \
        avl_tree_insert(map_entry_t(K, V))(map->tree, entry);                  \
    return &ref->value;                                                        \
  }                                                                            \
  bool map_erase(K, V)(map_t(K, V) * map, K key, V * out)                      \
  {                                                                            \
    map_entry_t(K, V) key_;                                                    \
    key_.key = key;                                                            \
                                                                               \
    map_entry_t(K, V) entry_out;                                               \
    bool ret = avl_tree_erase(map_entry_t(K, V))(map->tree, key_, &entry_out); \
    if (ret) {                                                                 \
      *out = entry_out.value;                                                  \
    }                                                                          \
    return ret;                                                                \
  }

#endif /* _SHARED_MAP_H_ */