// #NOTE: This Hash Table is borrowed from Jai's implementation! (With some tweaks)
// I made my own version that's arena-backed, but the mechanisms are the same. 

typedef u32 hash_result;

constexpr hash_result HASH_TABLE_FIRST_VALID_HASH = 2;
constexpr hash_result HASH_TABLE_REMOVED_HASH     = 1;

typedef hash_result (*Hash_Function)(void* key, s64 size);
typedef bool (*Hash_Compare_Function)(void* key1, void* key2);

template <typename T, typename U>
struct Table_Entry {
  using KeyType = T; using ValueType = U;
  T key;
  U value;
  hash_result hash;
};

template <typename T, typename U>
struct Table {
  using KeyType = T; using ValueType = U;
  Allocator allocator = {};
  ArrayView<Table_Entry<T, U>> entries = {};
  s64 allocated       = 0;
  s64 count           = 0;
  s64 slots_filled    = 0;

  Hash_Function         hash_function    = nullptr;
  Hash_Compare_Function compare_function = nullptr;
  
  u32 load_factor_percent = 70;
  
  s64 add_collisions  = 0;
  s64 find_collisions = 0;

  bool refill_removed = true;
  bool count_collisions = false;
};

template <typename T, typename U> bool table_is_valid (Table<T, U>* table) {
  if (table == nullptr) return false;
  if (table->entries.data == nullptr) return false;
  if (table->allocated == 0) return false;
  
  if (table->hash_function == nullptr) return false;
  if (table->compare_function == nullptr) return false;
  
  return true;
}

template <typename T, typename U> void table_init (Table<T, U>* table, s64 slots_to_allocate=64) {
  if (table->allocator.proc == nullptr) {
    table->allocator = context_allocator(); // #remember_allocator
  }
  push_allocator(table->allocator);
  
  s64 n = Next_Power_Of_Two(slots_to_allocate);
  
  table->entries = ArrayView<Table_Entry<T, U>>(n, false);
  table->allocated = n;
  
  for (s64 i = 0; i < n; i += 1) {
    table->entries[i].hash = 0;
  }
  // default hash and compare functions:
  // table->hash_function = table_hash_function_fnv1a;
  // table->compare_function = u64_keys_match;
}

// Adds given key value pair to the table, returns a pointer to the inserted value.
template <typename T, typename U> U* table_add (Table<T, U>* table, T key, U value) {
  Assert(table_is_valid(table));
  Assert(table->load_factor_percent < 100);
  
  if ( ((table->slots_filled + 1) * 100) >= (table->allocated * table->load_factor_percent) ) {
    table_resize(table, Next_Power_Of_Two(table->allocated + 64));
  }
  Assert(table->slots_filled < table->allocated); 
  
  // #Walk_Table
  u32 mask = (u32)(table->allocated - 1);
  
  u32 hash = table->hash_function(&key, sizeof(T));
  
  if (hash < HASH_TABLE_FIRST_VALID_HASH) {
    hash += HASH_TABLE_FIRST_VALID_HASH;
  }
  
  u32 index = hash & mask;
  
  u32 probe_increment = 1;
  u32 table_while_loop = table->entries[index].hash;
  while (table_while_loop) {
    if (table->refill_removed) {
      if (table->entries[index].hash == HASH_TABLE_REMOVED_HASH) {
        table->slots_filled -= 1; // 1 will get re-added below, for total increment 0.
        break;
      }
    }
    
    if (table->count_collisions) {
      table->add_collisions += 1;
    }
    
    index = (index + probe_increment) & mask;
    probe_increment += 1;
    
    table_while_loop = table->entries[index].hash;
  }
  
  // Walk_Table walked us to an unused entry, so add our new data into this slot:
  table->count        += 1;
  table->slots_filled += 1;
  
  Table_Entry<T, U>* entry = &table->entries[index];
  entry->hash  = hash;
  entry->key   = key;
  entry->value = value;
  
  return &entry->value;
}

template <typename T, typename U> U* table_find_pointer (Table<T, U>* table, T key) {
  Assert(table_is_valid(table));
  if (!table_is_valid(table)) return nullptr;
  
  // #Walk_Table
  u32 mask = (u32)(table->allocated - 1);
  
  u32 hash = table->hash_function(&key, sizeof(T));
  
  if (hash < HASH_TABLE_FIRST_VALID_HASH) {
    hash += HASH_TABLE_FIRST_VALID_HASH;
  }
  
  u32 index = hash & mask;
  
  u32 probe_increment = 1;
  u32 table_while_loop = table->entries[index].hash;
  while (table_while_loop) {
    Table_Entry<T, U>* entry = &table->entries[index];
    if (entry->hash == hash) {
      if (table->compare_function(&entry->key, &key)) {
        return &entry->value;
      }
    }

    if (table->count_collisions) { table->find_collisions += 1; }
    
    index = (index + probe_increment) & mask;
    probe_increment += 1;
    
    table_while_loop = table->entries[index].hash;
  }
  
  return nullptr;
}

template <typename T, typename U> U* table_set (Table<T, U>* table, T key, U value) {
  U* value_ptr = table_find_pointer(table, key);
  if (value_ptr) {
    (*value_ptr) = value;
    return value_ptr;
  } else {
    return table_add(table, key, value);
  }
}

template <typename T, typename U> void table_resize (Table<T, U>* table, s64 slots_to_allocate) {
  s64 initial_count = table->entries.count;
  Assert(slots_to_allocate > initial_count);
  if (slots_to_allocate <= initial_count) {
    return;
  }
  
  // #TODO: When you resize you need to reinsert all the values, so we 
  // need a temporary copy of the original data:
  ArrayView<Table_Entry<T, U>> old_entries = table->entries;
  
  s64 n = Next_Power_Of_Two(slots_to_allocate);
  
  table->entries = ArrayView<Table_Entry<T, U>>(n, false);
  table->allocated = n;
  
  table->count = 0;
  table->slots_filled = 0;
  
  // Initialize new values:
  for (s64 i = 0; i < old_entries.count; i += 1) {
    Table_Entry<T, U> entry = old_entries[i];
    // if entry is valid!
    if (entry.hash > HASH_TABLE_FIRST_VALID_HASH) {
      table_add(table, entry.key, entry.value);
    }
  }
}

template <typename T, typename U> void table_reset (Table<T, U>* table, bool keep_memory=true) {
  table->count = 0;
  table->slots_filled = 0;
  
  for (s64 i = 0; i < table->entries.count; i += 1) {
    table->entries[i].hash = 0;
  }
  
  if (!keep_memory) { array_reset(table->entries); }
}

template <typename T, typename U> void table_release (Table<T, U>* table) {
  arena_array_free(table->entries);
  
#if BUILD_DEBUG 
  poison_struct(table);
#endif
}


template <typename T, typename U> bool table_contains (Table<T, U>* table, T key) {
  return (table_find_pointer(table, key) != nullptr);
}

template <typename T, typename U> bool table_find (Table<T, U>* table, T key, U* value) {
  U* pointer = table_find_pointer(table, key);
  if (pointer) {
    (*value) = (*pointer);
    return true;
  }
  
  return false;
}

template <typename T, typename U> bool table_remove (Table<T, U>* table, T key, U* value) {
  Assert(table_is_valid(table));
  if (!table_is_valid(table)) return nullptr;
  
  // #Walk_Table
  u32 mask = (u32)(table->allocated - 1);
  
  u32 hash = table->hash_function(&key, sizeof(T));
  
  if (hash < HASH_TABLE_FIRST_VALID_HASH) {
    hash += HASH_TABLE_FIRST_VALID_HASH;
  }
  
  u32 index = hash & mask;
  
  u32 probe_increment = 1;
  u32 table_while_loop = table->entries[index].hash;
  while (table_while_loop) {
    Table_Entry<T, U>* entry = &table->entries[index];
    
    if ((entry->hash == hash) && table->compare_function(&entry->key, &key)) {
      entry->hash = HASH_TABLE_REMOVED_HASH;
      table->count -= 1;
      (*value) = entry->value;
      return true;
    }
    
    index = (index + probe_increment) & mask;
    probe_increment += 1;
    
    table_while_loop = table->entries[index].hash;
  }
  
  return false;
}

// #TODO: we need a for expansion iterator?
// table_find_multiple (put results in Temp-backed Array<>, and return it as an ArrayView<T>) {
template <typename T, typename U> ArrayView<U> table_find_multiple (Table<T, U>* table, T key, U* value) {
  Array<U> results;
  results.allocator = temp();
  
  // #Walk_Table
  u32 mask = (u32)(table->allocated - 1);
  
  u32 hash = table->hash_function(&key, sizeof(T));
  
  if (hash < HASH_TABLE_FIRST_VALID_HASH) {
    hash += HASH_TABLE_FIRST_VALID_HASH;
  }
  
  u32 index = hash & mask;
  
  u32 probe_increment = 1;
  u32 table_while_loop = table->entries[index].hash;
  while (table_while_loop) {
    Table_Entry<T, U>* entry = &table->entries[index];
    
    if (entry->hash == hash) {
      if (table->compare_function(&entry->key, &key)) {
        array_add(results, entry->value);
      } else {
        if (table->count_collisions) { table->find_collisions += 1; }
      }
    } else {
      if (table->count_collisions) { table->find_collisions += 1; }
    }
    
    index = (index + probe_increment) & mask;
    probe_increment += 1;
    
    table_while_loop = table->entries[index].hash;
  }
  
  return to_view(results);
}



// #TODO:
// find_or_add is kind of like table_set, but used when you
// just want a pointer to the value, which you can fill in.
template <typename T, typename U> U* table_find_or_add (Table<T, U>* table, T key, bool* newly_added) {
  U* value = table_find_pointer(table, key);
  if (value) {
    (*newly_added) = false;
    return value;
  }
  
  U new_value;
  value = table_add(table, key, new_value);
  (*newly_added) = true;
  return value;
}

// find_or_add :: (table: *Table, key: table.Key_Type) -> (entry: *table.Value_Type, newly_added: bool) {
//     value := table_find_pointer(table, key);
//     if value return value, false;

//     new_value: table.Value_Type;
//     value = table_add(table, key, new_value);
//     return value, true;
// }