// SPDX-License-Identifier: GPL-3.0-or-later

#include <errno.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "hashmap.h"

#define INITIAL_BUCKETS 64
#define LOAD_FACTOR_NUM 3
#define LOAD_FACTOR_DEN 4

typedef struct {
  char *key;
  void *value;
  bool occupied;
  size_t order;
} Bucket;

typedef struct {
  size_t bucket_idx;
} OrderEntry;

struct OrderedHashmap {
  Bucket *buckets;
  size_t n_buckets;
  OrderEntry *order;
  size_t n_entries;
  size_t n_order_alloc;
};

static size_t hash_string(const char *s) {
  size_t h = 5381;
  for (; *s; s++)
    h = h * 33 + (unsigned char)*s;
  return h;
}

static size_t probe(size_t hash, size_t i, size_t n) {
  return (hash + i) & (n - 1);
}

static bool is_power_of_2(size_t n) { return n && !(n & (n - 1)); }

static size_t next_power_of_2(size_t n) {
  size_t v = 1;
  while (v < n)
    v <<= 1;
  return v;
}

OrderedHashmap *ordered_hashmap_new(void) {
  OrderedHashmap *h = calloc(1, sizeof(*h));
  if (!h)
    return NULL;

  h->n_buckets = INITIAL_BUCKETS;
  h->buckets = calloc(h->n_buckets, sizeof(Bucket));
  if (!h->buckets) {
    free(h);
    return NULL;
  }

  h->n_order_alloc = 16;
  h->order = calloc(h->n_order_alloc, sizeof(OrderEntry));
  if (!h->order) {
    free(h->buckets);
    free(h);
    return NULL;
  }

  return h;
}

void ordered_hashmap_free(OrderedHashmap *h) {
  if (!h)
    return;
  for (size_t i = 0; i < h->n_buckets; i++) {
    if (h->buckets[i].occupied)
      free(h->buckets[i].key);
  }
  free(h->buckets);
  free(h->order);
  free(h);
}

void ordered_hashmap_freep(OrderedHashmap **hp) { ordered_hashmap_free(*hp); }

static ssize_t find_bucket(OrderedHashmap *h, const char *key) {
  size_t hash = hash_string(key);
  for (size_t i = 0; i < h->n_buckets; i++) {
    size_t idx = probe(hash, i, h->n_buckets);
    if (!h->buckets[idx].occupied)
      return -(ssize_t)idx - 1;
    if (strcmp(h->buckets[idx].key, key) == 0)
      return (ssize_t)idx;
  }
  return -(ssize_t)h->n_buckets - 1;
}

static int resize(OrderedHashmap *h) {
  size_t new_n = h->n_buckets * 2;
  if (!is_power_of_2(new_n))
    new_n = next_power_of_2(new_n);

  Bucket *new_buckets = calloc(new_n, sizeof(Bucket));
  if (!new_buckets)
    return -ENOMEM;

  for (size_t oi = 0; oi < h->n_entries; oi++) {
    Bucket *old = &h->buckets[h->order[oi].bucket_idx];
    size_t hash = hash_string(old->key);

    for (size_t i = 0;; i++) {
      size_t idx = probe(hash, i, new_n);
      if (!new_buckets[idx].occupied) {
        new_buckets[idx].key = old->key;
        new_buckets[idx].value = old->value;
        new_buckets[idx].occupied = true;
        new_buckets[idx].order = oi;
        h->order[oi].bucket_idx = idx;
        break;
      }
    }
  }

  free(h->buckets);
  h->buckets = new_buckets;
  h->n_buckets = new_n;
  return 0;
}

int ordered_hashmap_put(OrderedHashmap *h, const char *key, void *value) {
  if (h->n_entries * LOAD_FACTOR_DEN >= h->n_buckets * LOAD_FACTOR_NUM) {
    int r = resize(h);
    if (r < 0)
      return r;
  }

  ssize_t pos = find_bucket(h, key);
  if (pos >= 0)
    return -EEXIST;

  size_t idx = (size_t)(-(pos + 1));
  char *kcopy = strdup(key);
  if (!kcopy)
    return -ENOMEM;

  if (h->n_entries >= h->n_order_alloc) {
    size_t new_alloc = h->n_order_alloc * 2;
    OrderEntry *no = realloc(h->order, new_alloc * sizeof(OrderEntry));
    if (!no) {
      free(kcopy);
      return -ENOMEM;
    }
    h->order = no;
    h->n_order_alloc = new_alloc;
  }

  h->buckets[idx].key = kcopy;
  h->buckets[idx].value = value;
  h->buckets[idx].occupied = true;
  h->buckets[idx].order = h->n_entries;

  h->order[h->n_entries].bucket_idx = idx;
  h->n_entries++;

  return 0;
}

void *ordered_hashmap_get(OrderedHashmap *h, const char *key) {
  if (!h)
    return NULL;
  ssize_t pos = find_bucket(h, key);
  if (pos < 0)
    return NULL;
  return h->buckets[pos].value;
}

bool ordered_hashmap_contains(OrderedHashmap *h, const char *key) {
  if (!h)
    return false;
  return find_bucket(h, key) >= 0;
}

void *ordered_hashmap_remove(OrderedHashmap *h, const char *key) {
  if (!h)
    return NULL;

  ssize_t pos = find_bucket(h, key);
  if (pos < 0)
    return NULL;

  Bucket *b = &h->buckets[pos];
  void *val = b->value;
  size_t removed_order = b->order;

  free(b->key);
  b->key = NULL;
  b->value = NULL;
  b->occupied = false;

  for (size_t i = 1;; i++) {
    size_t ci = probe((size_t)pos, i, h->n_buckets);
    if (!h->buckets[ci].occupied)
      break;

    size_t natural = hash_string(h->buckets[ci].key) & (h->n_buckets - 1);
    bool needs_move;
    if ((size_t)pos < ci)
      needs_move = (natural <= (size_t)pos || natural > ci);
    else
      needs_move = (natural <= (size_t)pos && natural > ci);

    if (needs_move) {
      h->buckets[pos] = h->buckets[ci];
      h->order[h->buckets[pos].order].bucket_idx = (size_t)pos;
      h->buckets[ci].occupied = false;
      h->buckets[ci].key = NULL;
      h->buckets[ci].value = NULL;
      pos = (ssize_t)ci;
    }
  }

  for (size_t i = removed_order; i + 1 < h->n_entries; i++) {
    h->order[i] = h->order[i + 1];
    h->buckets[h->order[i].bucket_idx].order = i;
  }
  h->n_entries--;

  return val;
}

size_t ordered_hashmap_size(OrderedHashmap *h) { return h ? h->n_entries : 0; }

bool ordered_hashmap_iterate(OrderedHashmapIter *iter, const char **key,
                             void **value) {
  if (!iter->h || iter->idx >= iter->h->n_entries)
    return false;

  Bucket *b = &iter->h->buckets[iter->h->order[iter->idx].bucket_idx];
  if (key)
    *key = b->key;
  if (value)
    *value = b->value;
  iter->idx++;
  return true;
}
