#include "domain/search.h"

#include <stdint.h>
#include <stdlib.h>
#include <string.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 count;
  size_t capacity;
  size_t message_count;
};

static bool
decode_codepoints(const char *text, uint32_t **output, size_t *count)
{
  const size_t bytes = strlen(text);
  size_t offset = 0, used = 0;
  uint32_t *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;
    }
  return hash;
}

static gram_bucket_t *
find_bucket(const tiny_search_index_t *index, const uint32_t *points, unsigned width)
{
  const uint64_t hash = gram_hash(points, width);
  size_t position;
  for (position = 0; position < index->count; ++position)
    {
      if (index->buckets[position].hash == hash && index->buckets[position].width == width
          && memcmp(index->buckets[position].points, points, width * sizeof(uint32_t)) == 0)
        {
          return &index->buckets[position];
        }
    }
  return NULL;
}

static bool
add_gram(tiny_search_index_t *index, const uint32_t *points, unsigned width, size_t message_index)
{
  gram_bucket_t *bucket = find_bucket(index, points, width);
  if (bucket == NULL)
    {
      gram_bucket_t *resized;
      size_t capacity;
      if (index->count == index->capacity)
        {
          capacity = index->capacity ? index->capacity * 2 : 64;
          resized = realloc(index->buckets, capacity * sizeof(*resized));
          if (!resized)
            {
              return false;
            }
          index->buckets = resized;
          index->capacity = capacity;
        }
      bucket = &index->buckets[index->count++];
      memset(bucket, 0, sizeof(*bucket));
      bucket->hash = gram_hash(points, width);
      bucket->width = width;
      memcpy(bucket->points, points, width * sizeof(uint32_t));
    }
  if (bucket->count > 0 && bucket->messages[bucket->count - 1] == message_index)
    {
      return true;
    }
  if (bucket->count == bucket->capacity)
    {
      const size_t capacity = bucket->capacity ? bucket->capacity * 2 : 8;
      size_t *resized = realloc(bucket->messages, capacity * sizeof(*resized));
      if (!resized)
        {
          return false;
        }
      bucket->messages = resized;
      bucket->capacity = capacity;
    }
  bucket->messages[bucket->count++] = message_index;
  return true;
}

static bool
index_text(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(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));
}

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

bool
tiny_search_index_add(tiny_search_index_t *index,
                      const tiny_message_t *message,
                      size_t message_index)
{
  size_t tag;
  if (!index_text(index, message->content, message_index))
    {
      return false;
    }
  for (tag = 0; tag < message->metadata.tag_count; ++tag)
    {
      if (!index_text(index, message->metadata.tags[tag], message_index))
        {
          return false;
        }
    }
  if (message_index >= index->message_count)
    {
      index->message_count = message_index + 1;
    }
  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;
  *output = NULL;
  *count = 0;
  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
    {
      gram_bucket_t *bucket;
      width = point_count >= 3 ? 3 : 2;
      for (offset = 0; offset + width <= point_count; ++offset)
        {
          if ((bucket = find_bucket(index, points + offset, (unsigned)width)) != NULL)
            {
              size_t item;
              for (item = 0; item < bucket->count; ++item)
                {
                  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;
}

double
tiny_search_score(const char *query, const tiny_message_t *message)
{
  uint32_t *query_points = NULL, *message_points = NULL;
  size_t query_count = 0, message_count = 0, width, query_offset, message_offset, matches = 0,
         total, tag;
  double score = 0.0;
  if (!decode_codepoints(query, &query_points, &query_count)
      || !decode_codepoints(message->content, &message_points, &message_count) || query_count == 0)
    {
      goto done;
    }
  width = query_count >= 3 ? 3 : (query_count >= 2 ? 2 : 1);
  total = query_count - width + 1;
  for (query_offset = 0; query_offset < total; ++query_offset)
    {
      bool found = false;
      for (message_offset = 0; message_offset + width <= message_count; ++message_offset)
        {
          if (memcmp(query_points + query_offset,
                     message_points + message_offset,
                     width * sizeof(uint32_t))
              == 0)
            {
              found = true;
              break;
            }
        }
      for (tag = 0; tag < message->metadata.tag_count && !found; ++tag)
        {
          uint32_t *tag_points = NULL;
          size_t tag_count = 0;
          if (!decode_codepoints(message->metadata.tags[tag], &tag_points, &tag_count))
            {
              continue;
            }
          for (message_offset = 0; message_offset + width <= tag_count; ++message_offset)
            {
              if (memcmp(query_points + query_offset,
                         tag_points + message_offset,
                         width * sizeof(uint32_t))
                  == 0)
                {
                  found = true;
                  break;
                }
            }
          free(tag_points);
        }
      if (found)
        {
          matches++;
        }
    }
  score = (double)matches / (double)total;
  if (query_count <= message_count)
    {
      for (message_offset = 0; message_offset + query_count <= message_count; ++message_offset)
        {
          if (memcmp(query_points, message_points + message_offset, query_count * sizeof(uint32_t))
              == 0)
            {
              score += 1.0;
              break;
            }
        }
    }
done:
  free(query_points);
  free(message_points);
  return score;
}