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

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

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

#include "base/allocator.h"

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

static double
score_component(const char *query, const tiny_message_t *message, int component)
{
  tiny_message_t candidate = *message;
  candidate.metadata.tags = NULL;
  candidate.metadata.tag_count = 0;
  candidate.metadata.summary = NULL;
  if (component == 1)
    {
      candidate.content = message->metadata.summary ? message->metadata.summary : "";
    }
  else if (component == 2)
    {
      candidate.content = "";
      candidate.metadata.tags = message->metadata.tags;
      candidate.metadata.tag_count = message->metadata.tag_count;
    }
  return tiny_search_score(query, &candidate);
}

static int
compare_results(const void *left, const void *right)
{
  const result_item_t *a = left;
  const result_item_t *b = right;
  if (a->score > b->score)
    {
      return -1;
    }
  if (a->score < b->score)
    {
      return 1;
    }
  if (a->entry->message.metadata.importance > b->entry->message.metadata.importance)
    {
      return -1;
    }
  if (a->entry->message.metadata.importance < b->entry->message.metadata.importance)
    {
      return 1;
    }
  if (a->entry->message.updated_at_ms > b->entry->message.updated_at_ms)
    {
      return -1;
    }
  if (a->entry->message.updated_at_ms < b->entry->message.updated_at_ms)
    {
      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 void
call_search_messages(const mcpc_tool_t *tool, mcpc_ucbr_t *result)
{
  tiny_mcp_tools_context_t *context = tiny_mcp_tools_context_current();
  char *query = tiny_mcp_tool_get_string(tool, "query", TINY_QUERY_MAX);
  char *session = tiny_mcp_tool_get_string(tool, "session_id", TINY_SESSION_ID_MAX);
  char *kind_text = tiny_mcp_tool_get_string(tool, "memory_kind", 24);
  char *tier_text = tiny_mcp_tool_get_string(tool, "memory_tier", 16);
  char *tag = tiny_mcp_tool_get_string(tool, "tag", TINY_TAG_MAX);
  int32_t requested_limit = 20;
  int32_t minimum_importance = 0;
  mcpc_errcode_t limit_status
      = mcpc_tool_get_tpropval_i32(tool, (const char8_t *)"limit", &requested_limit);
  mcpc_errcode_t importance_status
      = mcpc_tool_get_tpropval_i32(tool, (const char8_t *)"min_importance", &minimum_importance);
  tiny_memory_kind_t kind = TINY_KIND_CONVERSATION;
  tiny_memory_tier_t tier = TINY_TIER_SHORT;
  size_t limit, index, count = 0, candidate_count = 0, *candidates = NULL;
  result_item_t *items = NULL;
  tiny_json_buffer_t json;
  bool valid = context && query && query[0]
               && (limit_status == MCPC_EC_0 || limit_status == MCPC_EC_ANYDATA_NOTFOUND)
               && requested_limit >= 1 && requested_limit <= 100;
  valid = valid && minimum_importance >= 0 && minimum_importance <= 100
          && (importance_status == MCPC_EC_0 || importance_status == MCPC_EC_ANYDATA_NOTFOUND)
          && (!kind_text || tiny_kind_parse(kind_text, &kind))
          && (!tier_text || tiny_tier_parse(tier_text, &tier));
  if (!valid)
    {
      tiny_mcp_tool_add_error(result, "VALIDATION_ERROR", "query or limit is invalid");
    }
  else
    {
      limit = (size_t)requested_limit;
      if (!context->search_index)
        {
          context->search_index = tiny_mcp_tools_build_search_index(context->store);
        }
      if (!context->search_index
          || !tiny_search_index_candidates(
              context->search_index, query, &candidates, &candidate_count))
        {
          tiny_mcp_tool_add_error(result, "INTERNAL_ERROR", "search index allocation failed");
        }
      else
        {
          items = malloc((candidate_count + 1) * sizeof(*items));
          if (!items)
            {
              tiny_mcp_tool_add_error(result, "INTERNAL_ERROR", "out of memory");
            }
          else
            {
              for (index = 0; index < candidate_count; ++index)
                {
                  const tiny_store_entry_t *entry
                      = tiny_store_entry(context->store, candidates[index]);
                  double score;
                  if (!entry || entry->deleted
                      || (session && session[0] && strcmp(entry->message.session_id, session) != 0))
                    {
                      continue;
                    }
                  if ((kind_text && entry->message.metadata.memory_kind != kind)
                      || (tier_text && entry->message.metadata.memory_tier != tier)
                      || entry->message.metadata.importance < minimum_importance)
                    {
                      continue;
                    }
                  if (tag)
                    {
                      size_t tag_index;
                      for (tag_index = 0; tag_index < entry->message.metadata.tag_count;
                           ++tag_index)
                        {
                          if (strcmp(tag, entry->message.metadata.tags[tag_index]) == 0)
                            {
                              break;
                            }
                        }
                      if (tag_index == entry->message.metadata.tag_count)
                        {
                          continue;
                        }
                    }
                  score = tiny_search_score(query, &entry->message);
                  if (score > 0)
                    {
                      items[count].entry = entry;
                      items[count].score = score;
                      items[count].content_score = score_component(query, &entry->message, 0);
                      items[count].summary_score = score_component(query, &entry->message, 1);
                      items[count].tag_score = score_component(query, &entry->message, 2);
                      count++;
                    }
                }
              qsort(items, count, sizeof(*items), compare_results);
              if (count > limit)
                {
                  count = limit;
                }
              tiny_json_init(&json);
              tiny_json_raw(&json, "{\"results\":[");
              for (index = 0; index < count; ++index)
                {
                  if (index)
                    {
                      tiny_json_raw(&json, ",");
                    }
                  {
                    tiny_json_buffer_t item;
                    tiny_json_init(&item);
                    tiny_mcp_tool_append_message(&item, &items[index].entry->message);
                    if (item.length > 0)
                      {
                        item.length--;
                        item.data[item.length] = '\0';
                      }
                    tiny_json_raw(&item, ",\"score\":");
                    tiny_json_number(&item, items[index].score);
                    tiny_json_raw(&item, ",\"score_breakdown\":{\"content\":");
                    tiny_json_number(&item, items[index].content_score);
                    tiny_json_raw(&item, ",\"summary\":");
                    tiny_json_number(&item, items[index].summary_score);
                    tiny_json_raw(&item, ",\"tags\":");
                    tiny_json_number(&item, items[index].tag_score);
                    tiny_json_raw(&item, ",\"importance\":0},\"matched_fields\":[");
                    {
                      bool matched = false;
                      if (items[index].content_score > 0)
                        {
                          tiny_json_raw(&item, "\"content\"");
                          matched = true;
                        }
                      if (items[index].summary_score > 0)
                        {
                          tiny_json_raw(&item, matched ? ",\"summary\"" : "\"summary\"");
                          matched = true;
                        }
                      if (items[index].tag_score > 0)
                        {
                          tiny_json_raw(&item, matched ? ",\"tags\"" : "\"tags\"");
                        }
                    }
                    tiny_json_raw(&item, "],\"relation_count\":");
                    tiny_json_u64(&item, items[index].entry->message.metadata.relation_count);
                    tiny_json_raw(&item, "}");
                    tiny_json_raw_n(&json, item.data, item.length);
                    tiny_json_free(&item);
                  }
                }
              tiny_json_raw(&json, "]}");
              tiny_mcp_tool_add_json(result, &json);
              tiny_json_free(&json);
            }
        }
    }
  free(items);
  free(candidates);
  free(query);
  free(session);
  free(kind_text);
  free(tier_text);
  free(tag);
}

bool
tiny_mcp_register_search_messages(mcpc_server_t *server)
{
  mcpc_tool_t *tool
      = tiny_mcp_tool_new("search_messages",
                          "Search persisted message content and tags by text relevance. Use it to "
                          "recall matching information globally or within one session.",
                          call_search_messages);
  return tool
         && tiny_mcp_tool_add_property(
             tool,
             "query",
             "UTF-8 text to match against persisted message content and tags.",
             MCPC_U8STR,
             true)
         && tiny_mcp_tool_add_property(
             tool,
             "session_id",
             "Conversation session to search; omit to search across all sessions.",
             MCPC_U8STR,
             false)
         && tiny_mcp_tool_add_property(
             tool, "memory_kind", "Optional exact memory kind filter.", MCPC_U8STR, false)
         && tiny_mcp_tool_add_property(
             tool, "memory_tier", "Optional exact memory tier filter.", MCPC_U8STR, false)
         && tiny_mcp_tool_add_property(tool, "tag", "Optional exact tag filter.", MCPC_U8STR, false)
         && tiny_mcp_tool_add_property(tool,
                                       "min_importance",
                                       "Minimum importance from 0 to 100; defaults to 0.",
                                       MCPC_I32,
                                       false)
         && tiny_mcp_tool_add_property(
             tool,
             "limit",
             "Maximum ranked matches to return, from 1 to 100; defaults to 20.",
             MCPC_I32,
             false)
         && mcpc_server_add_tool(server, tool) == MCPC_EC_0;
}