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

#include "mcp/tools/common/common.h"
#include "mcp/tools/common/context.h"
#include "mjson.h"

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

#include "base/allocator.h"

static bool
metadata_has_only_known_fields(const char *json, int length)
{
  int offset = 0, key_offset, key_length, value_offset, value_length, value_type;
  while ((offset = mjson_next(json,
                              length,
                              offset,
                              &key_offset,
                              &key_length,
                              &value_offset,
                              &value_length,
                              &value_type))
         != 0)
    {
      char key[32];
      int decoded = mjson_get_string(json + key_offset, key_length, "$", key, sizeof(key));
      (void)value_offset;
      (void)value_length;
      (void)value_type;
      if (decoded <= 0 || decoded >= (int)sizeof(key))
        {
          return false;
        }
      key[decoded] = '\0';
      if (strcmp(key, "tool_call_id") != 0 && strcmp(key, "tags") != 0
          && strcmp(key, "memory_tier") != 0 && strcmp(key, "memory_kind") != 0
          && strcmp(key, "importance") != 0 && strcmp(key, "summary") != 0
          && strcmp(key, "relations") != 0 && strcmp(key, "source_message_ids") != 0)
        {
          return false;
        }
    }
  return true;
}

bool
tiny_mcp_parse_metadata(const char *json, tiny_message_t *message, tiny_metadata_mask_t *present)
{
  const char *tool_call = NULL, *tier = NULL, *tags = NULL, *kind = NULL, *summary = NULL;
  const char *relations = NULL, *sources = NULL;
  int tool_call_len = 0, tier_len = 0, tags_len = 0, kind_len = 0, summary_len = 0;
  int relations_len = 0, sources_len = 0;
  int offset = 0, key_offset, key_length, value_offset, value_length, value_type;
  int length = (int)strlen(json);
  int tool_call_type, tier_type, tags_type;
  int kind_type, summary_type, importance_type;
  int relations_type, sources_type;
  double importance = 50;
  tiny_metadata_mask_t fields = 0;
  if (strcmp(json, "{}") == 0)
    {
      return true;
    }
  if (mjson_find(json, length, "$", NULL, NULL) != MJSON_TOK_OBJECT
      || !metadata_has_only_known_fields(json, length))
    {
      return false;
    }
  tool_call_type = mjson_find(json, length, "$.tool_call_id", &tool_call, &tool_call_len);
  tier_type = mjson_find(json, length, "$.memory_tier", &tier, &tier_len);
  tags_type = mjson_find(json, length, "$.tags", &tags, &tags_len);
  kind_type = mjson_find(json, length, "$.memory_kind", &kind, &kind_len);
  summary_type = mjson_find(json, length, "$.summary", &summary, &summary_len);
  importance_type = mjson_find(json, length, "$.importance", NULL, NULL);
  relations_type = mjson_find(json, length, "$.relations", &relations, &relations_len);
  sources_type = mjson_find(json, length, "$.source_message_ids", &sources, &sources_len);
  if ((tool_call_type != MJSON_TOK_INVALID && tool_call_type != MJSON_TOK_STRING)
      || (tier_type != MJSON_TOK_INVALID && tier_type != MJSON_TOK_STRING)
      || (tags_type != MJSON_TOK_INVALID && tags_type != MJSON_TOK_ARRAY)
      || (kind_type != MJSON_TOK_INVALID && kind_type != MJSON_TOK_STRING)
      || (summary_type != MJSON_TOK_INVALID && summary_type != MJSON_TOK_STRING)
      || (importance_type != MJSON_TOK_INVALID && importance_type != MJSON_TOK_NUMBER)
      || (relations_type != MJSON_TOK_INVALID && relations_type != MJSON_TOK_ARRAY)
      || (sources_type != MJSON_TOK_INVALID && sources_type != MJSON_TOK_ARRAY))
    {
      return false;
    }
  fields |= tool_call_type != MJSON_TOK_INVALID ? TINY_METADATA_TOOL_CALL_ID : 0;
  fields |= tags_type != MJSON_TOK_INVALID ? TINY_METADATA_TAGS : 0;
  fields |= tier_type != MJSON_TOK_INVALID ? TINY_METADATA_TIER : 0;
  fields |= kind_type != MJSON_TOK_INVALID ? TINY_METADATA_KIND : 0;
  fields |= importance_type != MJSON_TOK_INVALID ? TINY_METADATA_IMPORTANCE : 0;
  fields |= summary_type != MJSON_TOK_INVALID ? TINY_METADATA_SUMMARY : 0;
  fields |= relations_type != MJSON_TOK_INVALID ? TINY_METADATA_RELATIONS : 0;
  fields |= sources_type != MJSON_TOK_INVALID ? TINY_METADATA_SOURCES : 0;
  if (relations_type == MJSON_TOK_ARRAY)
    {
      offset = 0;
      while ((offset = mjson_next(relations,
                                  relations_len,
                                  offset,
                                  &key_offset,
                                  &key_length,
                                  &value_offset,
                                  &value_length,
                                  &value_type))
             != 0)
        {
          tiny_relation_t relation;
          char type[24];
          int decoded_type, decoded_id;
          tiny_relation_t *resized;
          if (value_type != MJSON_TOK_OBJECT
              || message->metadata.relation_count >= TINY_RELATION_COUNT_MAX)
            {
              return false;
            }
          decoded_type = mjson_get_string(
              relations + value_offset, value_length, "$.type", type, sizeof(type));
          decoded_id = mjson_get_string(relations + value_offset,
                                        value_length,
                                        "$.message_id",
                                        relation.message_id,
                                        sizeof(relation.message_id));
          if (decoded_type <= 0 || decoded_type >= (int)sizeof(type)
              || decoded_id != TINY_MESSAGE_ID_LENGTH)
            {
              return false;
            }
          type[decoded_type] = '\0';
          relation.message_id[decoded_id] = '\0';
          if (!tiny_relation_type_parse(type, &relation.type)
              || !tiny_message_id_valid(relation.message_id))
            {
              return false;
            }
          resized = realloc(message->metadata.relations,
                            (message->metadata.relation_count + 1) * sizeof(*resized));
          if (!resized)
            {
              return false;
            }
          message->metadata.relations = resized;
          message->metadata.relations[message->metadata.relation_count++] = relation;
        }
    }
  if (sources_type == MJSON_TOK_ARRAY)
    {
      offset = 0;
      while ((offset = mjson_next(sources,
                                  sources_len,
                                  offset,
                                  &key_offset,
                                  &key_length,
                                  &value_offset,
                                  &value_length,
                                  &value_type))
             != 0)
        {
          char *id;
          char **resized;
          int decoded;
          if (value_type != MJSON_TOK_STRING
              || message->metadata.source_count >= TINY_SOURCE_COUNT_MAX)
            {
              return false;
            }
          id = malloc(TINY_MESSAGE_ID_LENGTH + 1);
          if (!id)
            {
              return false;
            }
          decoded = mjson_get_string(
              sources + value_offset, value_length, "$", id, TINY_MESSAGE_ID_LENGTH + 1);
          if (decoded != TINY_MESSAGE_ID_LENGTH)
            {
              free(id);
              return false;
            }
          id[decoded] = '\0';
          if (!tiny_message_id_valid(id))
            {
              free(id);
              return false;
            }
          resized = realloc(message->metadata.source_message_ids,
                            (message->metadata.source_count + 1) * sizeof(*resized));
          if (!resized)
            {
              free(id);
              return false;
            }
          message->metadata.source_message_ids = resized;
          message->metadata.source_message_ids[message->metadata.source_count++] = id;
        }
    }
  if (importance_type == MJSON_TOK_NUMBER
      && mjson_get_number(json, length, "$.importance", &importance) > 0)
    {
      if (importance < 0 || importance > 100 || importance != (double)(uint8_t)importance)
        {
          return false;
        }
      message->metadata.importance = (uint8_t)importance;
    }
  if (kind_type == MJSON_TOK_STRING)
    {
      char value[24];
      int decoded = mjson_get_string(json, length, "$.memory_kind", value, sizeof(value));
      if (decoded <= 0 || decoded >= (int)sizeof(value))
        {
          return false;
        }
      value[decoded] = '\0';
      if (!tiny_kind_parse(value, &message->metadata.memory_kind))
        {
          return false;
        }
    }
  if (summary_type == MJSON_TOK_STRING)
    {
      int decoded;
      message->metadata.summary = malloc(TINY_SUMMARY_MAX + 1);
      if (!message->metadata.summary)
        {
          return false;
        }
      decoded = mjson_get_string(
          json, length, "$.summary", message->metadata.summary, TINY_SUMMARY_MAX + 1);
      if (decoded < 0 || decoded > TINY_SUMMARY_MAX
          || (decoded && !tiny_utf8_valid(message->metadata.summary, (size_t)decoded)))
        {
          return false;
        }
      message->metadata.summary[decoded] = '\0';
      if (decoded == 0)
        {
          free(message->metadata.summary);
          message->metadata.summary = NULL;
        }
    }
  if (tool_call_type == MJSON_TOK_STRING)
    {
      int decoded;
      message->metadata.tool_call_id = malloc(TINY_TOOL_CALL_ID_MAX + 1);
      if (!message->metadata.tool_call_id)
        {
          return false;
        }
      decoded = mjson_get_string(json,
                                 length,
                                 "$.tool_call_id",
                                 message->metadata.tool_call_id,
                                 TINY_TOOL_CALL_ID_MAX + 1);
      if (decoded < 0 || decoded > TINY_TOOL_CALL_ID_MAX
          || !tiny_utf8_valid(message->metadata.tool_call_id, (size_t)decoded))
        {
          return false;
        }
      message->metadata.tool_call_id[decoded] = '\0';
    }
  if (tier_type == MJSON_TOK_STRING)
    {
      char tier_value[16];
      int decoded = mjson_get_string(json, length, "$.memory_tier", tier_value, sizeof(tier_value));
      if (decoded < 0 || decoded >= (int)sizeof(tier_value))
        {
          return false;
        }
      tier_value[decoded] = '\0';
      if (!tiny_tier_parse(tier_value, &message->metadata.memory_tier))
        {
          return false;
        }
    }
  if (tags_type == MJSON_TOK_ARRAY)
    {
      while ((offset = mjson_next(tags,
                                  tags_len,
                                  offset,
                                  &key_offset,
                                  &key_length,
                                  &value_offset,
                                  &value_length,
                                  &value_type))
             != 0)
        {
          char *tag;
          char **resized;
          int decoded;
          if (value_type != 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(tags + value_offset, value_length, "$", 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(*message->metadata.tags));
          if (!resized)
            {
              free(tag);
              return false;
            }
          message->metadata.tags = resized;
          message->metadata.tags[message->metadata.tag_count++] = tag;
        }
    }
  if (present)
    {
      *present = fields;
    }
  return tiny_metadata_valid(&message->metadata);
}

static void
call_write_message(const mcpc_tool_t *tool, mcpc_ucbr_t *result)
{
  tiny_mcp_tools_context_t *context = tiny_mcp_tools_context_current();
  tiny_message_t input, written;
  tiny_store_status_t status;
  char *role = NULL, *metadata = NULL;
  tiny_json_buffer_t json;
  tiny_message_init(&input);
  tiny_message_init(&written);
  input.session_id = tiny_mcp_tool_get_string(tool, "session_id", TINY_SESSION_ID_MAX);
  role = tiny_mcp_tool_get_string(tool, "role", 16);
  input.content = tiny_mcp_tool_get_string(tool, "content", TINY_CONTENT_MAX);
  metadata = tiny_mcp_tool_get_string(tool, "metadata_json", TINY_CONTENT_MAX);
  bool valid = context && input.session_id && input.session_id[0] && role && input.content
               && input.content[0] && tiny_role_parse(role, &input.role)
               && tiny_mcp_parse_metadata(metadata ? metadata : "{}", &input, NULL);
  if (!valid)
    {
      tiny_mcp_tool_add_error(result, "VALIDATION_ERROR", "invalid write_message arguments");
    }
  else
    {
      status = tiny_store_write(context->store, &input, &written);
      if (status != TINY_STORE_OK)
        {
          tiny_mcp_tool_add_error(
              result, tiny_store_status_name(status), "message was not written");
        }
      else
        {
          tiny_mcp_tools_invalidate_indexes(context);
          tiny_json_init(&json);
          tiny_json_raw(&json, "{\"status\":\"OK\",\"message\":");
          tiny_mcp_tool_append_message(&json, &written);
          tiny_json_raw(&json, "}");
          tiny_mcp_tool_add_json(result, &json);
          tiny_json_free(&json);
        }
    }
  free(role);
  free(metadata);
  tiny_message_free(&input);
  tiny_message_free(&written);
}

bool
tiny_mcp_register_write_message(mcpc_server_t *server)
{
  mcpc_tool_t *tool
      = tiny_mcp_tool_new("write_message",
                          "Append and durably persist a conversation message. Use it to store "
                          "information that should remain available after server restarts.",
                          call_write_message);
  return tool
         && tiny_mcp_tool_add_property(
             tool, "session_id", "Conversation session that owns the message.", MCPC_U8STR, true)
         && tiny_mcp_tool_add_property(
             tool, "role", "Message author role: user, assistant, or system.", MCPC_U8STR, true)
         && tiny_mcp_tool_add_property(
             tool, "content", "UTF-8 message text to persist.", MCPC_U8STR, true)
         && tiny_mcp_tool_add_property(tool,
                                       "metadata_json",
                                       "Optional JSON object string containing tool_call_id, tags, "
                                       "memory_tier, memory_kind, importance, summary, relations, "
                                       "and/or source_message_ids; "
                                       "defaults to {}.",
                                       MCPC_U8STR,
                                       false)
         && mcpc_server_add_tool(server, tool) == MCPC_EC_0;
}