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

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

#include "base/allocator.h"

#define EMPTY_SLOT SIZE_MAX

typedef struct memory_bucket
{
  uint64_t hash;
  char message_id[TINY_MESSAGE_ID_LENGTH + 1];
  size_t entry_index;
  tiny_relation_reference_t *incoming;
  size_t incoming_count;
  size_t incoming_capacity;
} memory_bucket_t;

struct tiny_memory_index
{
  memory_bucket_t *buckets;
  size_t count;
  size_t *slots;
  size_t slot_capacity;
};

static uint64_t
id_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
slot_start(uint64_t hash, size_t capacity)
{
  return (size_t)(hash ^ (hash >> 32)) & (capacity - 1);
}

static memory_bucket_t *
find_bucket(const tiny_memory_index_t *index, const char *message_id)
{
  uint64_t hash;
  size_t slot, bucket;
  if (!index || !message_id || !index->slot_capacity)
    {
      return NULL;
    }
  hash = id_hash(message_id);
  slot = slot_start(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].message_id, message_id) == 0)
        {
          return &index->buckets[bucket];
        }
      slot = (slot + 1) & (index->slot_capacity - 1);
    }
  return NULL;
}

static bool
add_incoming(memory_bucket_t *bucket, size_t source, tiny_relation_type_t type)
{
  tiny_relation_reference_t *resized;
  size_t capacity;
  if (bucket->incoming_count == bucket->incoming_capacity)
    {
      capacity = bucket->incoming_capacity ? bucket->incoming_capacity * 2 : 4;
      if (capacity < bucket->incoming_capacity
          || !(resized = realloc(bucket->incoming, capacity * sizeof(*resized))))
        {
          return false;
        }
      bucket->incoming = resized;
      bucket->incoming_capacity = capacity;
    }
  bucket->incoming[bucket->incoming_count].source_index = source;
  bucket->incoming[bucket->incoming_count].type = type;
  bucket->incoming_count++;
  return true;
}

tiny_memory_index_t *
tiny_memory_index_build(const tiny_store_t *store)
{
  tiny_memory_index_t *index;
  size_t count = tiny_store_entry_count(store), capacity = 64, position;
  if (!store || count > SIZE_MAX / sizeof(memory_bucket_t))
    {
      return NULL;
    }
  while (count > (capacity / 10) * 7)
    {
      if (capacity > SIZE_MAX / 2)
        {
          return NULL;
        }
      capacity *= 2;
    }
  index = calloc(1, sizeof(*index));
  if (!index || (count && !(index->buckets = calloc(count, sizeof(*index->buckets))))
      || !(index->slots = malloc(capacity * sizeof(*index->slots))))
    {
      tiny_memory_index_free(index);
      return NULL;
    }
  index->count = count;
  index->slot_capacity = capacity;
  for (position = 0; position < capacity; ++position)
    {
      index->slots[position] = EMPTY_SLOT;
    }
  for (position = 0; position < count; ++position)
    {
      const tiny_store_entry_t *entry = tiny_store_entry(store, position);
      memory_bucket_t *bucket = &index->buckets[position];
      size_t slot;
      memcpy(bucket->message_id, entry->message.message_id, sizeof(bucket->message_id));
      bucket->hash = id_hash(bucket->message_id);
      bucket->entry_index = position;
      slot = slot_start(bucket->hash, capacity);
      while (index->slots[slot] != EMPTY_SLOT)
        {
          slot = (slot + 1) & (capacity - 1);
        }
      index->slots[slot] = position;
    }
  for (position = 0; position < count; ++position)
    {
      const tiny_store_entry_t *entry = tiny_store_entry(store, position);
      size_t relation;
      if (entry->deleted)
        {
          continue;
        }
      for (relation = 0; relation < entry->message.metadata.relation_count; ++relation)
        {
          const tiny_relation_t *value = &entry->message.metadata.relations[relation];
          memory_bucket_t *target = find_bucket(index, value->message_id);
          if (!target)
            {
              /* The relation target is missing (likely deleted during rebuild); skip it. */
              continue;
            }
          if (!add_incoming(target, position, value->type))
            {
              tiny_memory_index_free(index);
              return NULL;
            }
        }
    }
  return index;
}

uint64_t
tiny_memory_index_memory_estimate(const tiny_memory_index_t *index)
{
  uint64_t total;
  size_t position;
  if (!index)
    {
      return 0;
    }
  total = sizeof(*index) + index->count * sizeof(*index->buckets)
          + index->slot_capacity * sizeof(*index->slots);
  for (position = 0; position < index->count; ++position)
    {
      total += index->buckets[position].incoming_capacity
               * sizeof(*index->buckets[position].incoming);
    }
  return total;
}

void
tiny_memory_index_free(tiny_memory_index_t *index)
{
  size_t position;
  if (!index)
    {
      return;
    }
  for (position = 0; position < index->count; ++position)
    {
      free(index->buckets[position].incoming);
    }
  free(index->buckets);
  free(index->slots);
  free(index);
}

bool
tiny_memory_index_find(const tiny_memory_index_t *index,
                       const char *message_id,
                       size_t *entry_index)
{
  memory_bucket_t *bucket;
  if (!entry_index)
    {
      return false;
    }
  *entry_index = 0;
  bucket = find_bucket(index, message_id);
  if (bucket)
    {
      *entry_index = bucket->entry_index;
    }
  return bucket != NULL;
}

bool
tiny_memory_index_incoming(const tiny_memory_index_t *index,
                           const char *message_id,
                           const tiny_relation_reference_t **references,
                           size_t *count)
{
  memory_bucket_t *bucket;
  if (!references || !count)
    {
      return false;
    }
  *references = NULL;
  *count = 0;
  bucket = find_bucket(index, message_id);
  if (bucket)
    {
      *references = bucket->incoming;
      *count = bucket->incoming_count;
    }
  return bucket != NULL;
}