#include "mcp/tools/common/history.h"

#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include "base/allocator.h"

typedef struct history_bucket
{
  uint64_t hash;
  char *session_id;
  size_t *entries;
  size_t count;
  size_t capacity;
} history_bucket_t;

struct tiny_history_index
{
  history_bucket_t *buckets;
  size_t count;
  size_t capacity;
  size_t *slots;
  size_t slot_capacity;
};

#define EMPTY_SLOT SIZE_MAX

static uint64_t
session_hash(const char *value)
{
  uint64_t hash = UINT64_C(1469598103934665603);
  while (*value)
    {
      hash ^= (unsigned char)*value++;
      hash *= UINT64_C(1099511628211);
    }
  return hash;
}

static size_t
start_slot(uint64_t hash, size_t capacity)
{
  return (size_t)(hash ^ (hash >> 32)) & (capacity - 1);
}

static history_bucket_t *
find_bucket(const tiny_history_index_t *index, const char *session_id)
{
  uint64_t hash;
  size_t slot, bucket;
  if (!index || !session_id || !index->slot_capacity)
    {
      return NULL;
    }
  hash = session_hash(session_id);
  slot = start_slot(hash, index->slot_capacity);
  while ((bucket = index->slots[slot]) != EMPTY_SLOT)
    {
      if (bucket >= index->count)
        {
          return NULL;
        }
      if (index->buckets[bucket].hash == hash
          && strcmp(index->buckets[bucket].session_id, session_id) == 0)
        {
          return &index->buckets[bucket];
        }
      slot = (slot + 1) & (index->slot_capacity - 1);
    }
  return NULL;
}

static void
insert_slot(size_t *slots, size_t capacity, const history_bucket_t *buckets, size_t bucket)
{
  size_t slot = start_slot(buckets[bucket].hash, capacity);
  while (slots[slot] != EMPTY_SLOT)
    {
      slot = (slot + 1) & (capacity - 1);
    }
  slots[slot] = bucket;
}

static bool
reserve_table(tiny_history_index_t *index, size_t required)
{
  history_bucket_t *buckets;
  size_t *slots, capacity, position;
  if (required > index->capacity)
    {
      capacity = index->capacity ? index->capacity : 32;
      while (capacity < required)
        {
          if (capacity > SIZE_MAX / 2)
            {
              return false;
            }
          capacity *= 2;
        }
      if (capacity > SIZE_MAX / sizeof(*buckets))
        {
          return false;
        }
      buckets = realloc(index->buckets, capacity * sizeof(*buckets));
      if (!buckets)
        {
          return false;
        }
      index->buckets = buckets;
      index->capacity = capacity;
    }
  capacity = index->slot_capacity ? index->slot_capacity : 64;
  while (required > (capacity / 10) * 7)
    {
      if (capacity > SIZE_MAX / 2)
        {
          return false;
        }
      capacity *= 2;
    }
  if (capacity == index->slot_capacity)
    {
      return true;
    }
  slots = malloc(capacity * sizeof(*slots));
  if (!slots)
    {
      return false;
    }
  for (position = 0; position < capacity; ++position)
    {
      slots[position] = EMPTY_SLOT;
    }
  for (position = 0; position < index->count; ++position)
    {
      insert_slot(slots, capacity, index->buckets, position);
    }
  free(index->slots);
  index->slots = slots;
  index->slot_capacity = capacity;
  return true;
}

static int
compare_entry(const tiny_store_t *store, size_t left, size_t right)
{
  const tiny_message_t *a = &tiny_store_entry(store, left)->message;
  const tiny_message_t *b = &tiny_store_entry(store, right)->message;
  if (a->timestamp_ms < b->timestamp_ms)
    {
      return -1;
    }
  if (a->timestamp_ms > b->timestamp_ms)
    {
      return 1;
    }
  return strcmp(a->message_id, b->message_id);
}

bool
tiny_history_index_add(tiny_history_index_t *index, const tiny_store_t *store, size_t entry_index)
{
  const tiny_store_entry_t *entry = tiny_store_entry(store, entry_index);
  history_bucket_t *bucket;
  size_t low, high, capacity;
  if (!index || !entry || !entry->message.session_id)
    {
      return false;
    }
  bucket = find_bucket(index, entry->message.session_id);
  if (!bucket)
    {
      char *session = NULL;
      size_t *entries = malloc(8 * sizeof(*entries));
      size_t length = strlen(entry->message.session_id);
      if (!entries || !(session = malloc(length + 1)) || index->count == SIZE_MAX
          || !reserve_table(index, index->count + 1))
        {
          free(entries);
          free(session);
          return false;
        }
      memcpy(session, entry->message.session_id, length + 1);
      bucket = &index->buckets[index->count];
      memset(bucket, 0, sizeof(*bucket));
      bucket->hash = session_hash(session);
      bucket->session_id = session;
      bucket->entries = entries;
      bucket->capacity = 8;
      insert_slot(index->slots, index->slot_capacity, index->buckets, index->count++);
    }
  if (bucket->count == bucket->capacity)
    {
      size_t *resized;
      if (bucket->capacity > SIZE_MAX / 2
          || (capacity = bucket->capacity * 2) > SIZE_MAX / sizeof(*resized)
          || !(resized = realloc(bucket->entries, capacity * sizeof(*resized))))
        {
          return false;
        }
      bucket->entries = resized;
      bucket->capacity = capacity;
    }
  low = 0;
  high = bucket->count;
  while (low < high)
    {
      size_t middle = low + (high - low) / 2;
      if (compare_entry(store, bucket->entries[middle], entry_index) <= 0)
        {
          low = middle + 1;
        }
      else
        {
          high = middle;
        }
    }
  memmove(bucket->entries + low + 1,
          bucket->entries + low,
          (bucket->count - low) * sizeof(*bucket->entries));
  bucket->entries[low] = entry_index;
  bucket->count++;
  return true;
}

tiny_history_index_t *
tiny_history_index_build(const tiny_store_t *store)
{
  tiny_history_index_t *index = calloc(1, sizeof(*index));
  size_t position;
  if (!index)
    {
      return NULL;
    }
  for (position = 0; position < tiny_store_entry_count(store); ++position)
    {
      if (!tiny_store_entry(store, position)->deleted
          && !tiny_history_index_add(index, store, position))
        {
          tiny_history_index_free(index);
          return NULL;
        }
    }
  return index;
}

void
tiny_history_index_free(tiny_history_index_t *index)
{
  size_t position;
  if (!index)
    {
      return;
    }
  for (position = 0; position < index->count; ++position)
    {
      free(index->buckets[position].session_id);
      free(index->buckets[position].entries);
    }
  free(index->buckets);
  free(index->slots);
  free(index);
}

bool
tiny_history_index_session(const tiny_history_index_t *index,
                           const char *session_id,
                           const size_t **entries,
                           size_t *count)
{
  history_bucket_t *bucket;
  if (!entries || !count)
    {
      return false;
    }
  *entries = NULL;
  *count = 0;
  bucket = find_bucket(index, session_id);
  if (bucket)
    {
      *entries = bucket->entries;
      *count = bucket->count;
    }
  return index && session_id;
}

size_t
tiny_history_index_session_count(const tiny_history_index_t *index)
{
  return index ? index->count : 0;
}

bool
tiny_history_index_session_at(const tiny_history_index_t *index,
                              size_t position,
                              const char **session_id,
                              const size_t **entries,
                              size_t *count)
{
  if (!index || position >= index->count || !session_id || !entries || !count)
    {
      return false;
    }
  *session_id = index->buckets[position].session_id;
  *entries = index->buckets[position].entries;
  *count = index->buckets[position].count;
  return true;
}