#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)
        {
          return false;
        }
    }
  return true;
}

static bool
parse_metadata(const char *json, tiny_message_t *message)
{
  const char *tool_call = NULL, *tier = NULL, *tags = NULL;
  int tool_call_len = 0, tier_len = 0, tags_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;
  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);
  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))
    {
      return false;
    }
  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;
        }
    }
  return true;
}

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)
               && parse_metadata(metadata ? metadata : "{}", &input);
  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
        {
          if (context->search_index
              && !tiny_search_index_add(
                  context->search_index, &written, tiny_store_entry_count(context->store) - 1))
            {
              tiny_search_index_free(context->search_index);
              context->search_index = NULL;
            }
          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 a conversation message", call_write_message);
  return tool && tiny_mcp_tool_add_property(tool, "session_id", "Session ID", MCPC_U8STR, true)
         && tiny_mcp_tool_add_property(tool, "role", "user, assistant, or system", MCPC_U8STR, true)
         && tiny_mcp_tool_add_property(tool, "content", "Message text", MCPC_U8STR, true)
         && tiny_mcp_tool_add_property(tool,
                                       "metadata_json",
                                       "Optional metadata object JSON; defaults to {}",
                                       MCPC_U8STR,
                                       false)
         && mcpc_server_add_tool(server, tool) == MCPC_EC_0;
}