#include "domain/search.h"

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

#include "base/allocator.h"

typedef struct gram_bucket
{
  uint64_t hash;
  uint32_t points[3];
  unsigned width;
  size_t *messages;
  size_t count;
  size_t capacity;
} gram_bucket_t;

struct tiny_search_index
{
  gram_bucket_t *buckets;
  size_t bucket_count;
  size_t bucket_capacity;
  size_t *slots;
  size_t slot_capacity;
  size_t message_count;
};

#define EMPTY_SLOT SIZE_MAX
#define INITIAL_BUCKET_CAPACITY 64U
#define INITIAL_SLOT_CAPACITY 128U

static bool
decode_codepoints(const char *text, uint32_t **output, size_t *count)
{
  size_t bytes;
  size_t offset = 0, used = 0;
  uint32_t *points;
  if (output == NULL || count == NULL)
    {
      return false;
    }
  *output = NULL;
  *count = 0;
  if (text == NULL)
    {
      return false;
    }
  bytes = strlen(text);
  if (!tiny_utf8_valid(text, bytes) || bytes > SIZE_MAX / sizeof(uint32_t) - 1)
    {
      return false;
    }
  points = malloc((bytes + 1) * sizeof(uint32_t));
  if (points == NULL)
    {
      return false;
    }
  while (offset < bytes)
    {
      unsigned char first = (unsigned char)text[offset++];
      uint32_t point;
      unsigned continuation;
      if (first < 0x80)
        {
          point = first;
          continuation = 0;
        }
      else if ((first & 0xe0) == 0xc0)
        {
          point = first & 0x1f;
          continuation = 1;
        }
      else if ((first & 0xf0) == 0xe0)
        {
          point = first & 0x0f;
          continuation = 2;
        }
      else
        {
          point = first & 0x07;
          continuation = 3;
        }
      while (continuation-- > 0)
        {
          point = (point << 6) | ((unsigned char)text[offset++] & 0x3f);
        }
      if (point >= 'A' && point <= 'Z')
        {
          point += 'a' - 'A';
        }
      points[used++] = point;
    }
  *output = points;
  *count = used;
  return true;
}

static uint64_t
gram_hash(const uint32_t *points, size_t width)
{
  uint64_t hash = 1469598103934665603ULL;
  size_t index;
  for (index = 0; index < width; ++index)
    {
      hash ^= points[index];
      hash *= 1099511628211ULL;
    }
#ifdef TINY_SEARCH_TEST_HASH_BITS
  hash &= (UINT64_C(1) << TINY_SEARCH_TEST_HASH_BITS) - 1;
#endif
  return hash;
}

static bool
bucket_matches(const gram_bucket_t *bucket, uint64_t hash, const uint32_t *points, unsigned width)
{
  return bucket->hash == hash && bucket->width == width
         && memcmp(bucket->points, points, width * sizeof(uint32_t)) == 0;
}

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

static void
insert_slot(size_t *slots, size_t capacity, const gram_bucket_t *buckets, size_t bucket_index)
{
  size_t position = slot_start(buckets[bucket_index].hash, capacity);
  while (slots[position] != EMPTY_SLOT)
    {
      position = (position + 1) & (capacity - 1);
    }
  slots[position] = bucket_index;
}

static const gram_bucket_t *
find_bucket_const(const tiny_search_index_t *index, const uint32_t *points, unsigned width)
{
  uint64_t hash;
  size_t position, bucket_index;
  if (index->slot_capacity == 0)
    {
      return NULL;
    }
  hash = gram_hash(points, width);
  position = slot_start(hash, index->slot_capacity);
  while ((bucket_index = index->slots[position]) != EMPTY_SLOT)
    {
      if (bucket_index >= index->bucket_count)
        {
          return NULL;
        }
      if (bucket_matches(&index->buckets[bucket_index], hash, points, width))
        {
          return &index->buckets[bucket_index];
        }
      position = (position + 1) & (index->slot_capacity - 1);
    }
  return NULL;
}

static gram_bucket_t *
find_bucket_mutable(tiny_search_index_t *index, const uint32_t *points, unsigned width)
{
  return (gram_bucket_t *)find_bucket_const(index, points, width);
}

static bool
reserve_buckets(tiny_search_index_t *index, size_t required)
{
  gram_bucket_t *resized;
  size_t capacity = index->bucket_capacity;
  if (required <= capacity)
    {
      return true;
    }
  if (required > SIZE_MAX / sizeof(*resized))
    {
      return false;
    }
  capacity = capacity ? capacity : INITIAL_BUCKET_CAPACITY;
  while (capacity < required)
    {
      if (capacity > SIZE_MAX / 2)
        {
          return false;
        }
      capacity *= 2;
    }
  resized = realloc(index->buckets, capacity * sizeof(*resized));
  if (!resized)
    {
      return false;
    }
  index->buckets = resized;
  index->bucket_capacity = capacity;
  return true;
}

static bool
ensure_slots(tiny_search_index_t *index, size_t required_buckets)
{
  size_t capacity = index->slot_capacity;
  size_t *slots;
  size_t position;
  if (capacity != 0 && required_buckets <= (capacity / 10) * 7)
    {
      return true;
    }
  capacity = capacity ? capacity : INITIAL_SLOT_CAPACITY;
  while (required_buckets > (capacity / 10) * 7)
    {
      if (capacity > SIZE_MAX / 2)
        {
          return false;
        }
      capacity *= 2;
    }
  if (capacity > SIZE_MAX / sizeof(*slots))
    {
      return false;
    }
  slots = malloc(capacity * sizeof(*slots));
  if (!slots)
    {
      return false;
    }
  for (position = 0; position < capacity; ++position)
    {
      slots[position] = EMPTY_SLOT;
    }
  for (position = 0; position < index->bucket_count; ++position)
    {
      insert_slot(slots, capacity, index->buckets, position);
    }
  free(index->slots);
  index->slots = slots;
  index->slot_capacity = capacity;
  return true;
}

static bool
reserve_messages(gram_bucket_t *bucket, size_t required)
{
  size_t capacity = bucket->capacity;
  size_t *resized;
  if (required <= capacity)
    {
      return true;
    }
  if (required > SIZE_MAX / sizeof(*resized))
    {
      return false;
    }
  capacity = capacity ? capacity : 8;
  while (capacity < required)
    {
      if (capacity > SIZE_MAX / 2)
        {
          return false;
        }
      capacity *= 2;
    }
  resized = realloc(bucket->messages, capacity * sizeof(*resized));
  if (!resized)
    {
      return false;
    }
  bucket->messages = resized;
  bucket->capacity = capacity;
  return true;
}

static bool
add_gram_direct(tiny_search_index_t *index,
                const uint32_t *points,
                unsigned width,
                size_t message_index)
{
  gram_bucket_t *bucket = find_bucket_mutable(index, points, width);
  if (bucket != NULL)
    {
      if (bucket->count > 0 && bucket->messages[bucket->count - 1] == message_index)
        {
          return true;
        }
      if (bucket->count == SIZE_MAX || !reserve_messages(bucket, bucket->count + 1))
        {
          return false;
        }
      bucket->messages[bucket->count++] = message_index;
      return true;
    }
  else
    {
      size_t *messages = malloc(8 * sizeof(*messages));
      size_t new_index;
      if (!messages)
        {
          return false;
        }
      if (index->bucket_count == SIZE_MAX || !reserve_buckets(index, index->bucket_count + 1)
          || !ensure_slots(index, index->bucket_count + 1))
        {
          free(messages);
          return false;
        }
      new_index = index->bucket_count;
      bucket = &index->buckets[new_index];
      memset(bucket, 0, sizeof(*bucket));
      bucket->hash = gram_hash(points, width);
      bucket->width = width;
      memcpy(bucket->points, points, width * sizeof(uint32_t));
      messages[0] = message_index;
      bucket->messages = messages;
      bucket->count = 1;
      bucket->capacity = 8;
      insert_slot(index->slots, index->slot_capacity, index->buckets, new_index);
      index->bucket_count++;
      return true;
    }
}

static bool
index_text_direct(tiny_search_index_t *index, const char *text, size_t message_index)
{
  uint32_t *points = NULL;
  size_t count = 0, offset;
  unsigned width;
  if (!decode_codepoints(text, &points, &count))
    {
      return false;
    }
  for (width = 2; width <= 3; ++width)
    {
      for (offset = 0; offset + width <= count; ++offset)
        {
          if (!add_gram_direct(index, points + offset, width, message_index))
            {
              free(points);
              return false;
            }
        }
    }
  free(points);
  return true;
}

tiny_search_index_t *
tiny_search_index_create(void)
{
  return calloc(1, sizeof(tiny_search_index_t));
}

uint64_t
tiny_search_index_memory_estimate(const tiny_search_index_t *index)
{
  uint64_t total;
  size_t position;
  if (!index)
    {
      return 0;
    }
  total = sizeof(*index) + index->bucket_capacity * sizeof(*index->buckets)
          + index->slot_capacity * sizeof(*index->slots);
  for (position = 0; position < index->bucket_count; ++position)
    {
      total += index->buckets[position].capacity * sizeof(*index->buckets[position].messages);
    }
  return total;
}

void
tiny_search_index_free(tiny_search_index_t *index)
{
  size_t position;
  if (!index)
    {
      return;
    }
  for (position = 0; position < index->bucket_count; ++position)
    {
      free(index->buckets[position].messages);
    }
  free(index->buckets);
  free(index->slots);
  free(index);
}

static bool
merge_delta(tiny_search_index_t *index, tiny_search_index_t *delta, size_t message_index)
{
  size_t position, new_buckets = 0, required;
  for (position = 0; position < delta->bucket_count; ++position)
    {
      gram_bucket_t *source = &delta->buckets[position];
      if (!find_bucket_const(index, source->points, source->width))
        {
          if (new_buckets == SIZE_MAX)
            {
              return false;
            }
          new_buckets++;
        }
    }
  if (new_buckets > SIZE_MAX - index->bucket_count)
    {
      return false;
    }
  required = index->bucket_count + new_buckets;
  if (!reserve_buckets(index, required) || !ensure_slots(index, required))
    {
      return false;
    }
  for (position = 0; position < delta->bucket_count; ++position)
    {
      gram_bucket_t *source = &delta->buckets[position];
      gram_bucket_t *target = find_bucket_mutable(index, source->points, source->width);
      if (target && !(target->count > 0 && target->messages[target->count - 1] == message_index)
          && (target->count == SIZE_MAX || !reserve_messages(target, target->count + 1)))
        {
          return false;
        }
    }
  for (position = 0; position < delta->bucket_count; ++position)
    {
      gram_bucket_t *source = &delta->buckets[position];
      gram_bucket_t *target = find_bucket_mutable(index, source->points, source->width);
      if (target)
        {
          if (!(target->count > 0 && target->messages[target->count - 1] == message_index))
            {
              target->messages[target->count++] = message_index;
            }
        }
      else
        {
          size_t new_index = index->bucket_count;
          index->buckets[new_index] = *source;
          source->messages = NULL;
          source->count = 0;
          source->capacity = 0;
          insert_slot(index->slots, index->slot_capacity, index->buckets, new_index);
          index->bucket_count++;
        }
    }
  index->message_count
      = message_index + 1 > index->message_count ? message_index + 1 : index->message_count;
  return true;
}

bool
tiny_search_index_add(tiny_search_index_t *index,
                      const tiny_message_t *message,
                      size_t message_index)
{
  tiny_search_index_t *delta;
  size_t tag;
  if (index == NULL || message == NULL || message->content == NULL
      || (message->metadata.tag_count != 0 && message->metadata.tags == NULL)
      || message_index == SIZE_MAX)
    {
      return false;
    }
  delta = tiny_search_index_create();
  if (!delta)
    {
      return false;
    }
  if (!index_text_direct(delta, message->content, message_index))
    {
      tiny_search_index_free(delta);
      return false;
    }
  if (message->metadata.summary
      && !index_text_direct(delta, message->metadata.summary, message_index))
    {
      tiny_search_index_free(delta);
      return false;
    }
  for (tag = 0; tag < message->metadata.tag_count; ++tag)
    {
      if (!index_text_direct(delta, message->metadata.tags[tag], message_index))
        {
          tiny_search_index_free(delta);
          return false;
        }
    }
  if (!merge_delta(index, delta, message_index))
    {
      tiny_search_index_free(delta);
      return false;
    }
  tiny_search_index_free(delta);
  return true;
}

bool
tiny_search_index_candidates(const tiny_search_index_t *index,
                             const char *query,
                             size_t **output,
                             size_t *count)
{
  uint32_t *points = NULL;
  size_t point_count = 0, width, offset, used = 0;
  bool *seen;
  size_t *items;
  if (output == NULL || count == NULL)
    {
      return false;
    }
  *output = NULL;
  *count = 0;
  if (index == NULL || query == NULL || index->message_count > SIZE_MAX / sizeof(*items))
    {
      return false;
    }
  if (!decode_codepoints(query, &points, &point_count))
    {
      return false;
    }
  items = malloc((index->message_count ? index->message_count : 1) * sizeof(*items));
  seen = calloc(index->message_count ? index->message_count : 1, sizeof(*seen));
  if (!items || !seen)
    {
      free(points);
      free(items);
      free(seen);
      return false;
    }
  if (point_count < 2)
    {
      for (offset = 0; offset < index->message_count; ++offset)
        {
          seen[offset] = true;
          items[used++] = offset;
        }
    }
  else
    {
      const gram_bucket_t *bucket;
      width = point_count >= 3 ? 3 : 2;
      for (offset = 0; offset + width <= point_count; ++offset)
        {
          if ((bucket = find_bucket_const(index, points + offset, (unsigned)width)) != NULL)
            {
              size_t item;
              for (item = 0; item < bucket->count; ++item)
                {
                  if (bucket->messages[item] >= index->message_count)
                    {
                      free(points);
                      free(items);
                      free(seen);
                      return false;
                    }
                  if (!seen[bucket->messages[item]])
                    {
                      seen[bucket->messages[item]] = true;
                      items[used++] = bucket->messages[item];
                    }
                }
            }
        }
    }
  free(points);
  free(seen);
  *output = items;
  *count = used;
  return true;
}

static bool
points_contain(const uint32_t *haystack,
               size_t haystack_count,
               const uint32_t *needle,
               size_t width)
{
  size_t offset;
  for (offset = 0; offset + width <= haystack_count; ++offset)
    {
      if (memcmp(haystack + offset, needle, width * sizeof(uint32_t)) == 0)
        {
          return true;
        }
    }
  return false;
}

static bool
score_points_breakdown(const uint32_t *query_points,
                       size_t query_count,
                       const uint32_t *content_points,
                       size_t content_count,
                       const tiny_message_t *message,
                       tiny_search_breakdown_t *breakdown)
{
  size_t width, query_offset, total;
  bool *content_hits = NULL, *summary_hits = NULL, *tag_hits = NULL;
  double exact_bonus = 0.0;
  width = query_count >= 3 ? 3 : (query_count >= 2 ? 2 : 1);
  total = query_count - width + 1;
  content_hits = calloc(total, sizeof(bool));
  summary_hits = calloc(total, sizeof(bool));
  tag_hits = calloc(total, sizeof(bool));
  if (!content_hits || !summary_hits || !tag_hits)
    {
      free(content_hits);
      free(summary_hits);
      free(tag_hits);
      return false;
    }
  for (query_offset = 0; query_offset < total; ++query_offset)
    {
      if (points_contain(content_points, content_count, query_points + query_offset, width))
        {
          content_hits[query_offset] = true;
        }
      else
        {
          size_t tag;
          for (tag = 0; tag < message->metadata.tag_count; ++tag)
            {
              uint32_t *tag_points = NULL;
              size_t tag_count = 0;
              if (!decode_codepoints(message->metadata.tags[tag], &tag_points, &tag_count))
                {
                  continue;
                }
              if (points_contain(tag_points, tag_count, query_points + query_offset, width))
                {
                  tag_hits[query_offset] = true;
                  free(tag_points);
                  break;
                }
              free(tag_points);
            }
          if (!tag_hits[query_offset] && message->metadata.summary)
            {
              uint32_t *summary_points = NULL;
              size_t summary_count = 0;
              if (decode_codepoints(message->metadata.summary, &summary_points, &summary_count)
                  && points_contain(
                      summary_points, summary_count, query_points + query_offset, width))
                {
                  summary_hits[query_offset] = true;
                }
              free(summary_points);
            }
        }
    }
  if (query_count <= content_count)
    {
      for (query_offset = 0; query_offset + query_count <= content_count; ++query_offset)
        {
          if (memcmp(query_points, content_points + query_offset, query_count * sizeof(uint32_t))
              == 0)
            {
              exact_bonus = 1.0;
              break;
            }
        }
    }
  {
    size_t total_matches = 0, content_matches = 0, summary_matches = 0, tag_matches = 0;
    for (query_offset = 0; query_offset < total; ++query_offset)
      {
        bool any
            = content_hits[query_offset] || summary_hits[query_offset] || tag_hits[query_offset];
        total_matches += any ? 1 : 0;
        content_matches += content_hits[query_offset] ? 1 : 0;
        summary_matches += summary_hits[query_offset] ? 1 : 0;
        tag_matches += tag_hits[query_offset] ? 1 : 0;
      }
    breakdown->total = total == 0 ? 0.0 : ((double)total_matches / (double)total) + exact_bonus;
    breakdown->content = total == 0 ? 0.0 : ((double)content_matches / (double)total) + exact_bonus;
    breakdown->summary = total == 0 ? 0.0 : (double)summary_matches / (double)total;
    breakdown->tags = total == 0 ? 0.0 : (double)tag_matches / (double)total;
  }
  free(content_hits);
  free(summary_hits);
  free(tag_hits);
  return true;
}

bool
tiny_search_score_breakdown(const char *query,
                            const tiny_message_t *message,
                            tiny_search_breakdown_t *breakdown)
{
  uint32_t *query_points = NULL, *content_points = NULL;
  size_t query_count = 0, content_count = 0;
  bool result;
  if (!breakdown)
    {
      return false;
    }
  memset(breakdown, 0, sizeof(*breakdown));
  if (query == NULL || message == NULL || message->content == NULL
      || (message->metadata.tag_count != 0 && message->metadata.tags == NULL))
    {
      return true;
    }
  if (!decode_codepoints(query, &query_points, &query_count) || query_count == 0
      || !decode_codepoints(message->content, &content_points, &content_count))
    {
      free(query_points);
      free(content_points);
      return false;
    }
  result = score_points_breakdown(
      query_points, query_count, content_points, content_count, message, breakdown);
  free(query_points);
  free(content_points);
  return result;
}

double
tiny_search_score(const char *query, const tiny_message_t *message)
{
  tiny_search_breakdown_t breakdown;
  if (!tiny_search_score_breakdown(query, message, &breakdown))
    {
      return 0.0;
    }
  return breakdown.total;
}