#include "mcp/server.h"

#include "domain/search.h"
#include "mcp/json.h"
#include "mjson.h"

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

#define MCP_PROTOCOL_VERSION "2024-11-05"

typedef struct result_item
{
  const tiny_store_entry_t *entry;
  double score;
} result_item_t;

static char *
read_line(FILE *input)
{
  size_t length = 0, capacity = 4096;
  int character;
  char *line = malloc(capacity);
  if (!line)
    {
      return NULL;
    }
  while ((character = fgetc(input)) != EOF)
    {
      if (length + 1 >= capacity)
        {
          char *r;
          capacity *= 2;
          r = realloc(line, capacity);
          if (!r)
            {
              free(line);
              return NULL;
            }
          line = r;
        }
      line[length++] = (char)character;
      if (character == '\n')
        {
          break;
        }
    }
  if (length == 0 && character == EOF)
    {
      free(line);
      return NULL;
    }
  line[length] = '\0';
  return line;
}

static char *
json_string_value(
    const char *json, int length, const char *path, size_t maximum, bool required, bool *valid)
{
  const char *token = NULL;
  int token_length = 0;
  int type = mjson_find(json, length, path, &token, &token_length);
  char *value;
  int decoded;
  if (type == MJSON_TOK_INVALID)
    {
      if (required)
        {
          *valid = false;
        }
      return NULL;
    }
  if (type != MJSON_TOK_STRING || token_length < 2 || (size_t)token_length > maximum * 6 + 2)
    {
      *valid = false;
      return NULL;
    }
  value = malloc(maximum + 1);
  if (!value)
    {
      *valid = false;
      return NULL;
    }
  decoded = mjson_get_string(json, length, path, value, (int)maximum + 1);
  if (decoded < 0 || (size_t)decoded > maximum || !tiny_utf8_valid(value, (size_t)decoded))
    {
      free(value);
      *valid = false;
      return NULL;
    }
  value[decoded] = '\0';
  return value;
}

static bool
object_has_only(const char *json, int length, const char *const *allowed, size_t allowed_count)
{
  int offset = 0, koff, klen, voff, vlen, vtype;
  while ((offset = mjson_next(json, length, offset, &koff, &klen, &voff, &vlen, &vtype)) != 0)
    {
      char key[64];
      int decoded;
      size_t index;
      if (klen <= 0)
        {
          return false;
        }
      decoded = mjson_get_string(json + koff, klen, "$", key, sizeof(key));
      if (decoded <= 0 || decoded >= (int)sizeof(key))
        {
          return false;
        }
      key[decoded] = '\0';
      for (index = 0; index < allowed_count; ++index)
        {
          if (strcmp(key, allowed[index]) == 0)
            {
              break;
            }
        }
      if (index == allowed_count)
        {
          return false;
        }
    }
  return true;
}

static void
free_input_message(tiny_message_t *message)
{
  tiny_message_free(message);
}

static bool
parse_tags(const char *json, int length, tiny_message_t *message)
{
  const char *array = NULL;
  int array_length = 0;
  int type = mjson_find(json, length, "$.metadata.tags", &array, &array_length);
  int offset = 0, koff, klen, voff, vlen, vtype;
  if (type == MJSON_TOK_INVALID)
    {
      return true;
    }
  if (type != MJSON_TOK_ARRAY)
    {
      return false;
    }
  while ((offset = mjson_next(array, array_length, offset, &koff, &klen, &voff, &vlen, &vtype))
         != 0)
    {
      char *tag;
      int decoded;
      char **resized;
      if (vtype != MJSON_TOK_STRING || message->metadata.tag_count >= TINY_TAG_COUNT_MAX)
        {
          return false;
        }
      tag = malloc(TINY_TAG_MAX + 1);
      if (!tag)
        {
          return false;
        }
      decoded = mjson_get_string(array + voff, vlen, "$", tag, TINY_TAG_MAX + 1);
      if (decoded <= 0 || decoded > TINY_TAG_MAX || !tiny_utf8_valid(tag, (size_t)decoded))
        {
          free(tag);
          return false;
        }
      tag[decoded] = '\0';
      resized = realloc(message->metadata.tags, (message->metadata.tag_count + 1) * sizeof(char *));
      if (!resized)
        {
          free(tag);
          return false;
        }
      message->metadata.tags = resized;
      message->metadata.tags[message->metadata.tag_count++] = tag;
    }
  return true;
}

static bool
parse_write(const char *args, int length, tiny_message_t *message)
{
  static const char *const top_keys[] = { "session_id", "role", "content", "metadata" };
  static const char *const metadata_keys[] = { "tool_call_id", "tags", "memory_tier" };
  bool valid = true;
  char *role = NULL, *tier = NULL;
  const char *metadata = NULL;
  int metadata_length = 0, metadata_type;
  tiny_message_init(message);
  if (!object_has_only(args, length, top_keys, sizeof(top_keys) / sizeof(top_keys[0])))
    {
      valid = false;
    }
  metadata_type = mjson_find(args, length, "$.metadata", &metadata, &metadata_length);
  if (metadata_type != MJSON_TOK_INVALID
      && (metadata_type != MJSON_TOK_OBJECT
          || !object_has_only(metadata,
                              metadata_length,
                              metadata_keys,
                              sizeof(metadata_keys) / sizeof(metadata_keys[0]))))
    {
      valid = false;
    }
  message->session_id
      = json_string_value(args, length, "$.session_id", TINY_SESSION_ID_MAX, true, &valid);
  role = json_string_value(args, length, "$.role", 16, true, &valid);
  message->content = json_string_value(args, length, "$.content", TINY_CONTENT_MAX, true, &valid);
  message->metadata.tool_call_id = json_string_value(
      args, length, "$.metadata.tool_call_id", TINY_TOOL_CALL_ID_MAX, false, &valid);
  tier = json_string_value(args, length, "$.metadata.memory_tier", 16, false, &valid);
  if (!valid || !message->session_id || !message->content || !role || message->session_id[0] == '\0'
      || message->content[0] == '\0' || !tiny_role_parse(role, &message->role)
      || (tier && !tiny_tier_parse(tier, &message->metadata.memory_tier))
      || !parse_tags(args, length, message))
    {
      valid = false;
    }
  free(role);
  free(tier);
  if (!valid)
    {
      free_input_message(message);
    }
  return valid;
}

static void
append_message(tiny_json_buffer_t *json, const tiny_message_t *message)
{
  size_t index;
  tiny_json_raw(json, "{\"message_id\":");
  tiny_json_string(json, message->message_id);
  tiny_json_raw(json, ",\"session_id\":");
  tiny_json_string(json, message->session_id);
  tiny_json_raw(json, ",\"timestamp\":");
  tiny_json_i64(json, message->timestamp_ms);
  tiny_json_raw(json, ",\"role\":");
  tiny_json_string(json, tiny_role_name(message->role));
  tiny_json_raw(json, ",\"content\":");
  tiny_json_string(json, message->content);
  tiny_json_raw(json, ",\"metadata\":{");
  if (message->metadata.tool_call_id)
    {
      tiny_json_raw(json, "\"tool_call_id\":");
      tiny_json_string(json, message->metadata.tool_call_id);
      tiny_json_raw(json, ",");
    }
  tiny_json_raw(json, "\"tags\":[");
  for (index = 0; index < message->metadata.tag_count; ++index)
    {
      if (index)
        {
          tiny_json_raw(json, ",");
        }
      tiny_json_string(json, message->metadata.tags[index]);
    }
  tiny_json_raw(json, "],\"memory_tier\":");
  tiny_json_string(json, tiny_tier_name(message->metadata.memory_tier));
  tiny_json_raw(json, "}}");
}

static void
send_json_result(
    FILE *output, const char *id, int id_length, tiny_json_buffer_t *inner, bool is_error)
{
  tiny_json_buffer_t response;
  tiny_json_init(&response);
  tiny_json_raw(&response, "{\"jsonrpc\":\"2.0\",\"id\":");
  tiny_json_raw_n(&response, id, (size_t)id_length);
  tiny_json_raw(&response, ",\"result\":{\"content\":[{\"type\":\"text\",\"text\":");
  tiny_json_string(&response, inner->data ? inner->data : "{}");
  tiny_json_raw(&response, "}]");
  if (is_error)
    {
      tiny_json_raw(&response, ",\"isError\":true");
    }
  tiny_json_raw(&response, "}}\n");
  if (!response.failed)
    {
      fwrite(response.data, 1, response.length, output);
      fflush(output);
    }
  tiny_json_free(&response);
}

static void
send_error(FILE *output, const char *id, int id_length, const char *code, const char *message)
{
  tiny_json_buffer_t inner;
  tiny_json_init(&inner);
  tiny_json_raw(&inner, "{\"status\":\"ERROR\",\"error\":{\"code\":");
  tiny_json_string(&inner, code);
  tiny_json_raw(&inner, ",\"message\":");
  tiny_json_string(&inner, message);
  tiny_json_raw(&inner, "}}");
  send_json_result(output, id, id_length, &inner, true);
  tiny_json_free(&inner);
}

static int
compare_history(const void *left, const void *right)
{
  const tiny_store_entry_t *a = *(const tiny_store_entry_t *const *)left,
                           *b = *(const tiny_store_entry_t *const *)right;
  if (a->message.timestamp_ms < b->message.timestamp_ms)
    {
      return -1;
    }
  if (a->message.timestamp_ms > b->message.timestamp_ms)
    {
      return 1;
    }
  return strcmp(a->message.message_id, b->message.message_id);
}

static int
compare_results(const void *l, const void *r)
{
  const result_item_t *a = l, *b = r;
  if (a->score > b->score)
    {
      return -1;
    }
  if (a->score < b->score)
    {
      return 1;
    }
  if (a->entry->message.timestamp_ms > b->entry->message.timestamp_ms)
    {
      return -1;
    }
  if (a->entry->message.timestamp_ms < b->entry->message.timestamp_ms)
    {
      return 1;
    }
  return strcmp(a->entry->message.message_id, b->entry->message.message_id);
}

static tiny_search_index_t *
build_search_index(const tiny_store_t *store)
{
  tiny_search_index_t *index = tiny_search_index_create();
  size_t position;
  if (!index)
    {
      return NULL;
    }
  for (position = 0; position < tiny_store_entry_count(store); ++position)
    {
      if (!tiny_search_index_add(index, &tiny_store_entry(store, position)->message, position))
        {
          tiny_search_index_free(index);
          return NULL;
        }
    }
  return index;
}

static void
handle_call(FILE *output,
            const char *id,
            int id_length,
            const char *name,
            const char *args,
            int args_length,
            tiny_store_t *store,
            tiny_search_index_t **search_index)
{
  tiny_json_buffer_t result;
  tiny_store_status_t status;
  tiny_json_init(&result);
  if (strcmp(name, "write_message") == 0)
    {
      tiny_message_t input, written;
      if (!parse_write(args, args_length, &input))
        {
          send_error(output, id, id_length, "VALIDATION_ERROR", "invalid write_message arguments");
          return;
        }
      tiny_message_init(&written);
      status = tiny_store_write(store, &input, &written);
      tiny_message_free(&input);
      if (status != TINY_STORE_OK)
        {
          send_error(
              output, id, id_length, tiny_store_status_name(status), "message was not written");
          return;
        }
      if (*search_index
          && !tiny_search_index_add(*search_index, &written, tiny_store_entry_count(store) - 1))
        {
          tiny_search_index_free(*search_index);
          *search_index = NULL;
        }
      tiny_json_raw(&result, "{\"status\":\"OK\",\"message\":");
      append_message(&result, &written);
      tiny_json_raw(&result, "}");
      tiny_message_free(&written);
    }
  else if (strcmp(name, "read_messages") == 0)
    {
      static const char *const keys[] = { "session_id" };
      bool valid = object_has_only(args, args_length, keys, 1);
      char *session
          = json_string_value(args, args_length, "$.session_id", TINY_SESSION_ID_MAX, true, &valid);
      size_t i, count = 0;
      const tiny_store_entry_t **items;
      if (!valid || !session || !session[0])
        {
          free(session);
          send_error(output, id, id_length, "VALIDATION_ERROR", "session_id is required");
          return;
        }
      items = malloc(tiny_store_entry_count(store) * sizeof(*items));
      if (!items)
        {
          free(session);
          send_error(output, id, id_length, "INTERNAL_ERROR", "out of memory");
          return;
        }
      for (i = 0; i < tiny_store_entry_count(store); ++i)
        {
          const tiny_store_entry_t *e = tiny_store_entry(store, i);
          if (!e->deleted && strcmp(e->message.session_id, session) == 0)
            {
              items[count++] = e;
            }
        }
      qsort(items, count, sizeof(*items), compare_history);
      tiny_json_raw(&result, "{\"history\":[");
      for (i = 0; i < count; ++i)
        {
          if (i)
            {
              tiny_json_raw(&result, ",");
            }
          append_message(&result, &items[i]->message);
        }
      tiny_json_raw(&result, "]}");
      free(items);
      free(session);
    }
  else if (strcmp(name, "search_messages") == 0)
    {
      static const char *const keys[] = { "query", "session_id", "limit" };
      bool valid = object_has_only(args, args_length, keys, 3);
      char *query = json_string_value(args, args_length, "$.query", TINY_QUERY_MAX, true, &valid);
      char *session = json_string_value(
          args, args_length, "$.session_id", TINY_SESSION_ID_MAX, false, &valid);
      double number = 20;
      size_t limit = 20, i, count = 0, candidate_count = 0, *candidates = NULL;
      result_item_t *items;
      if (mjson_find(args, args_length, "$.limit", NULL, NULL) != MJSON_TOK_INVALID)
        {
          if (!mjson_get_number(args, args_length, "$.limit", &number) || number < 1 || number > 100
              || floor(number) != number)
            {
              valid = false;
            }
          else
            {
              limit = (size_t)number;
            }
        }
      if (!valid || !query || !query[0])
        {
          free(query);
          free(session);
          send_error(output, id, id_length, "VALIDATION_ERROR", "query or limit is invalid");
          return;
        }
      if (!*search_index)
        {
          *search_index = build_search_index(store);
        }
      if (!*search_index
          || !tiny_search_index_candidates(*search_index, query, &candidates, &candidate_count))
        {
          free(query);
          free(session);
          send_error(output, id, id_length, "INTERNAL_ERROR", "search index allocation failed");
          return;
        }
      items = malloc((candidate_count ? candidate_count : 1) * sizeof(*items));
      if (!items)
        {
          free(candidates);
          free(query);
          free(session);
          send_error(output, id, id_length, "INTERNAL_ERROR", "out of memory");
          return;
        }
      for (i = 0; i < candidate_count; ++i)
        {
          const tiny_store_entry_t *e = tiny_store_entry(store, candidates[i]);
          double score;
          if (!e || e->deleted || (session && strcmp(e->message.session_id, session) != 0))
            {
              continue;
            }
          score = tiny_search_score(query, &e->message);
          if (score > 0)
            {
              items[count].entry = e;
              items[count].score = score;
              count++;
            }
        }
      qsort(items, count, sizeof(*items), compare_results);
      if (count > limit)
        {
          count = limit;
        }
      tiny_json_raw(&result, "{\"results\":[");
      for (i = 0; i < count; ++i)
        {
          if (i)
            {
              tiny_json_raw(&result, ",");
            }
          append_message(&result, &items[i].entry->message);
        }
      tiny_json_raw(&result, "]}");
      free(candidates);
      free(items);
      free(query);
      free(session);
    }
  else if (strcmp(name, "delete_message") == 0)
    {
      static const char *const keys[] = { "message_id", "session_id" };
      bool valid = object_has_only(args, args_length, keys, 2);
      char *mid = json_string_value(
          args, args_length, "$.message_id", TINY_MESSAGE_ID_LENGTH, true, &valid);
      char *session = json_string_value(
          args, args_length, "$.session_id", TINY_SESSION_ID_MAX, false, &valid);
      if (!valid || !mid || !tiny_message_id_valid(mid))
        {
          free(mid);
          free(session);
          send_error(output,
                     id,
                     id_length,
                     "VALIDATION_ERROR",
                     "message_id must be a canonical 26-character ULID");
          return;
        }
      status = tiny_store_delete(store, mid, session);
      free(mid);
      free(session);
      if (status != TINY_STORE_OK && status != TINY_STORE_NOT_FOUND)
        {
          send_error(output,
                     id,
                     id_length,
                     tiny_store_status_name(status),
                     "delete record was not written");
          return;
        }
      tiny_json_raw(&result,
                    status == TINY_STORE_OK ? "{\"status\":\"OK\"}" : "{\"status\":\"NOT_FOUND\"}");
    }
  else if (strcmp(name, "rebuild_index") == 0)
    {
      tiny_rebuild_report_t report;
      tiny_search_index_t *rebuilt;
      if (!object_has_only(args, args_length, NULL, 0))
        {
          send_error(output, id, id_length, "VALIDATION_ERROR", "rebuild_index takes no arguments");
          return;
        }
      status = tiny_store_rebuild(store, &report);
      if (status != TINY_STORE_OK)
        {
          send_error(output, id, id_length, tiny_store_status_name(status), "index rebuild failed");
          return;
        }
      rebuilt = build_search_index(store);
      if (!rebuilt)
        {
          tiny_search_index_free(*search_index);
          *search_index = NULL;
          send_error(output,
                     id,
                     id_length,
                     "INTERNAL_ERROR",
                     "log was compacted but the search index could not be allocated");
          return;
        }
      tiny_search_index_free(*search_index);
      *search_index = rebuilt;
      tiny_json_raw(&result, "{\"status\":\"OK\",\"messages_indexed\":");
      tiny_json_u64(&result, report.messages_indexed);
      tiny_json_raw(&result, ",\"records_removed\":");
      tiny_json_u64(&result, report.records_removed);
      tiny_json_raw(&result, ",\"bytes_before\":");
      tiny_json_u64(&result, report.bytes_before);
      tiny_json_raw(&result, ",\"bytes_after\":");
      tiny_json_u64(&result, report.bytes_after);
      tiny_json_raw(&result, "}");
    }
  else if (strcmp(name, "get_system_health") == 0)
    {
      tiny_store_health_t health;
      if (!object_has_only(args, args_length, NULL, 0))
        {
          send_error(
              output, id, id_length, "VALIDATION_ERROR", "get_system_health takes no arguments");
          return;
        }
      tiny_store_health(store, &health);
      tiny_json_raw(&result, "{\"status\":");
      tiny_json_string(&result, health.last_error[0] ? "DEGRADED" : "OK");
      tiny_json_raw(&result, ",\"format_version\":1,\"log_size\":");
      tiny_json_u64(&result, health.log_size);
      tiny_json_raw(&result, ",\"snapshot_size\":");
      tiny_json_u64(&result, health.snapshot_size);
      tiny_json_raw(&result, ",\"active_messages\":");
      tiny_json_u64(&result, health.active_messages);
      tiny_json_raw(&result, ",\"deleted_messages\":");
      tiny_json_u64(&result, health.deleted_messages);
      tiny_json_raw(&result, ",\"last_sequence\":");
      tiny_json_u64(&result, health.last_sequence);
      tiny_json_raw(&result, ",\"recovered_tail\":");
      tiny_json_raw(&result, health.recovered_tail ? "true" : "false");
      tiny_json_raw(&result, ",\"snapshot_valid\":");
      tiny_json_raw(&result, health.snapshot_valid ? "true" : "false");
      tiny_json_raw(&result, ",\"last_error\":");
      tiny_json_string(&result, health.last_error);
      tiny_json_raw(&result, "}");
    }
  else
    {
      send_error(output, id, id_length, "VALIDATION_ERROR", "unknown tool");
      return;
    }
  if (result.failed)
    {
      send_error(output, id, id_length, "INTERNAL_ERROR", "response allocation failed");
    }
  else
    {
      send_json_result(output, id, id_length, &result, false);
    }
  tiny_json_free(&result);
}

static const char *tools_json
    = "{\"tools\":["
      "{\"name\":\"write_message\",\"description\":\"Durably append a conversation "
      "message\",\"inputSchema\":{\"type\":\"object\",\"properties\":{\"session_id\":{\"type\":"
      "\"string\",\"maxLength\":256},\"role\":{\"type\":\"string\",\"enum\":[\"user\","
      "\"assistant\",\"system\"]},\"content\":{\"type\":\"string\",\"maxLength\":1048576},"
      "\"metadata\":{\"type\":\"object\",\"properties\":{\"tool_call_id\":{\"type\":\"string\","
      "\"maxLength\":256},\"tags\":{\"type\":\"array\",\"maxItems\":32,\"items\":{\"type\":"
      "\"string\",\"maxLength\":128}},\"memory_tier\":{\"type\":\"string\",\"enum\":[\"short\","
      "\"mid\",\"long\",\"discard\"]}},\"additionalProperties\":false}},\"required\":[\"session_"
      "id\",\"role\",\"content\"],\"additionalProperties\":false}},"
      "{\"name\":\"read_messages\",\"description\":\"Read a session in chronological "
      "order\",\"inputSchema\":{\"type\":\"object\",\"properties\":{\"session_id\":{\"type\":"
      "\"string\",\"maxLength\":256}},\"required\":[\"session_id\"],\"additionalProperties\":false}"
      "},"
      "{\"name\":\"search_messages\",\"description\":\"Search message content and "
      "tags\",\"inputSchema\":{\"type\":\"object\",\"properties\":{\"query\":{\"type\":\"string\","
      "\"maxLength\":4096},\"session_id\":{\"type\":\"string\",\"maxLength\":256},\"limit\":{"
      "\"type\":\"integer\",\"minimum\":1,\"maximum\":100}},\"required\":[\"query\"],"
      "\"additionalProperties\":false}},"
      "{\"name\":\"delete_message\",\"description\":\"Durably tombstone a "
      "message\",\"inputSchema\":{\"type\":\"object\",\"properties\":{\"message_id\":{\"type\":"
      "\"string\",\"minLength\":26,\"maxLength\":26},\"session_id\":{\"type\":\"string\","
      "\"maxLength\":256}},\"required\":[\"message_id\"],\"additionalProperties\":false}},"
      "{\"name\":\"rebuild_index\",\"description\":\"Compact the log and rebuild the "
      "snapshot\",\"inputSchema\":{\"type\":\"object\",\"properties\":{},\"additionalProperties\":"
      "false}},"
      "{\"name\":\"get_system_health\",\"description\":\"Report persistence health and "
      "counters\",\"inputSchema\":{\"type\":\"object\",\"properties\":{},\"additionalProperties\":"
      "false}}]}";

int
tiny_convmem_server_run(FILE *input, FILE *output, tiny_store_t *store)
{
  char *line;
  tiny_search_index_t *search_index = build_search_index(store);
  if (!search_index)
    {
      return 1;
    }
  while ((line = read_line(input)) != NULL)
    {
      int length = (int)strlen(line);
      const char *method = NULL, *id = NULL, *params = NULL;
      int ml = 0, il = 0, pl = 0;
      if (mjson_find(line, length, "$.method", &method, &ml) != MJSON_TOK_STRING)
        {
          free(line);
          continue;
        }
      mjson_find(line, length, "$.id", &id, &il);
      mjson_find(line, length, "$.params", &params, &pl);
      {
        char method_value[64];
        int decoded
            = mjson_get_string(line, length, "$.method", method_value, sizeof(method_value));
        if (decoded < 0)
          {
            free(line);
            continue;
          }
        method_value[decoded] = '\0';
        if (strcmp(method_value, "initialize") == 0 && id)
          {
            fprintf(output,
                    "{\"jsonrpc\":\"2.0\",\"id\":%.*s,\"result\":{\"protocolVersion\":\"%s\","
                    "\"capabilities\":{\"tools\":{}},\"serverInfo\":{\"name\":\"tiny-convmem-mcp\","
                    "\"version\":\"0.1.0\"}}}\n",
                    il,
                    id,
                    MCP_PROTOCOL_VERSION);
          }
        else if (strcmp(method_value, "tools/list") == 0 && id)
          {
            fprintf(
                output, "{\"jsonrpc\":\"2.0\",\"id\":%.*s,\"result\":%s}\n", il, id, tools_json);
          }
        else if (strcmp(method_value, "tools/call") == 0 && id && params)
          {
            char name[64];
            const char *args = NULL;
            int al = 0;
            int n = mjson_get_string(params, pl, "$.name", name, sizeof(name));
            if (n > 0)
              {
                name[n] = '\0';
                if (mjson_find(params, pl, "$.arguments", &args, &al) == MJSON_TOK_OBJECT)
                  {
                    handle_call(output, id, il, name, args, al, store, &search_index);
                  }
                else
                  {
                    send_error(output, id, il, "VALIDATION_ERROR", "arguments must be an object");
                  }
              }
          }
        fflush(output);
      }
      free(line);
    }
  tiny_search_index_free(search_index);
  return 0;
}