#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"
#include "mjson.h"

#define SEARCH_RESPONSE_BUDGET (8U * 1024U * 1024U - 64U * 1024U)
#define SEARCH_MAX_LIMIT 100
#define SEARCH_MAX_CONTEXT_SIZE 3

typedef struct search_filters
{
  bool kinds[8];
  bool tiers[8];
  char **tags;
  size_t tag_count;
} search_filters_t;

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 void
search_filters_free(search_filters_t *filters)
{
  size_t index;
  if (!filters)
    {
      return;
    }
  for (index = 0; index < filters->tag_count; ++index)
    {
      free(filters->tags[index]);
    }
  free(filters->tags);
  filters->tags = NULL;
  filters->tag_count = 0;
}

static bool
parse_filters_json(const char *json, search_filters_t *filters)
{
  int offset = 0, key_offset, key_length, value_offset, value_length, value_type;
  if (!json || !json[0])
    {
      return true;
    }
  if (mjson_find(json, (int)strlen(json), "$", NULL, NULL) != MJSON_TOK_OBJECT)
    {
      return false;
    }
  while ((offset = mjson_next(json,
                              (int)strlen(json),
                              offset,
                              &key_offset,
                              &key_length,
                              &value_offset,
                              &value_length,
                              &value_type))
         != 0)
    {
      char key[16];
      int decoded = mjson_get_string(json + key_offset, key_length, "$", key, sizeof(key));
      if (decoded <= 0 || decoded >= (int)sizeof(key))
        {
          return false;
        }
      key[decoded] = '\0';
      if (value_type != MJSON_TOK_ARRAY)
        {
          return false;
        }
      if (strcmp(key, "kinds") == 0)
        {
          int element = 0, e_key, e_klen, e_val, e_vlen, e_type;
          while ((element = mjson_next(json + value_offset,
                                       value_length,
                                       element,
                                       &e_key,
                                       &e_klen,
                                       &e_val,
                                       &e_vlen,
                                       &e_type))
                 != 0)
            {
              char name[24];
              tiny_memory_kind_t kind;
              int name_len
                  = mjson_get_string(json + value_offset + e_val, e_vlen, "$", name, sizeof(name));
              if (name_len <= 0 || name_len >= (int)sizeof(name) || !tiny_kind_parse(name, &kind))
                {
                  return false;
                }
              filters->kinds[kind] = true;
            }
        }
      else if (strcmp(key, "tiers") == 0)
        {
          int element = 0, e_key, e_klen, e_val, e_vlen, e_type;
          while ((element = mjson_next(json + value_offset,
                                       value_length,
                                       element,
                                       &e_key,
                                       &e_klen,
                                       &e_val,
                                       &e_vlen,
                                       &e_type))
                 != 0)
            {
              char name[16];
              tiny_memory_tier_t tier;
              int name_len
                  = mjson_get_string(json + value_offset + e_val, e_vlen, "$", name, sizeof(name));
              if (name_len <= 0 || name_len >= (int)sizeof(name) || !tiny_tier_parse(name, &tier))
                {
                  return false;
                }
              filters->tiers[tier] = true;
            }
        }
      else if (strcmp(key, "tags") == 0)
        {
          int element = 0, e_key, e_klen, e_val, e_vlen, e_type;
          while ((element = mjson_next(json + value_offset,
                                       value_length,
                                       element,
                                       &e_key,
                                       &e_klen,
                                       &e_val,
                                       &e_vlen,
                                       &e_type))
                 != 0)
            {
              char *tag;
              char name[TINY_TAG_MAX + 1];
              int name_len
                  = mjson_get_string(json + value_offset + e_val, e_vlen, "$", name, sizeof(name));
              if (name_len <= 0 || name_len >= (int)sizeof(name) || !name[0]
                  || strlen(name) > TINY_TAG_MAX)
                {
                  return false;
                }
              tag = malloc(name_len + 1);
              if (!tag)
                {
                  return false;
                }
              memcpy(tag, name, name_len + 1);
              if (filters->tag_count == SIZE_MAX)
                {
                  free(tag);
                  return false;
                }
              {
                char **resized = realloc(filters->tags, (filters->tag_count + 1) * sizeof(char *));
                if (!resized)
                  {
                    free(tag);
                    return false;
                  }
                filters->tags = resized;
                filters->tags[filters->tag_count++] = tag;
              }
            }
        }
      else
        {
          return false;
        }
    }
  return true;
}

static bool
filters_match(const search_filters_t *filters, const tiny_message_t *message)
{
  size_t index, tag_index;
  if (filters->kinds[TINY_KIND_CONVERSATION] || filters->kinds[TINY_KIND_EPISODE]
      || filters->kinds[TINY_KIND_SEMANTIC] || filters->kinds[TINY_KIND_PROCEDURE]
      || filters->kinds[TINY_KIND_PREFERENCE] || filters->kinds[TINY_KIND_SUMMARY])
    {
      if (!filters->kinds[message->metadata.memory_kind])
        {
          return false;
        }
    }
  if (filters->tiers[TINY_TIER_SHORT] || filters->tiers[TINY_TIER_MID]
      || filters->tiers[TINY_TIER_LONG] || filters->tiers[TINY_TIER_DISCARD])
    {
      if (!filters->tiers[message->metadata.memory_tier])
        {
          return false;
        }
    }
  for (index = 0; index < filters->tag_count; ++index)
    {
      for (tag_index = 0; tag_index < message->metadata.tag_count; ++tag_index)
        {
          if (strcmp(filters->tags[index], message->metadata.tags[tag_index]) == 0)
            {
              break;
            }
        }
      if (tag_index == message->metadata.tag_count)
        {
          return false;
        }
    }
  return true;
}

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
append_context_preview(tiny_json_buffer_t *json,
                       tiny_mcp_tools_context_t *context,
                       const tiny_message_t *message,
                       int32_t context_size)
{
  const size_t *entries = NULL;
  size_t entry_count = 0, target_position = SIZE_MAX, position;
  int32_t before = context_size, after = context_size;
  tiny_json_raw(json, ",\"context_preview\":{\"before\":[");
  if (!tiny_history_index_session(
          context->history_index, message->session_id, &entries, &entry_count))
    {
      before = 0;
      after = 0;
    }
  else
    {
      for (position = 0; position < entry_count; ++position)
        {
          const tiny_store_entry_t *entry = tiny_store_entry(context->store, entries[position]);
          if (entry && !entry->deleted
              && strcmp(entry->message.message_id, message->message_id) == 0)
            {
              target_position = position;
              break;
            }
        }
    }
  if (target_position == SIZE_MAX)
    {
      before = 0;
      after = 0;
    }
  for (position = (target_position > (size_t)before ? target_position - (size_t)before : 0);
       position < target_position;
       ++position)
    {
      const tiny_store_entry_t *entry = tiny_store_entry(context->store, entries[position]);
      tiny_json_buffer_t item;
      if (!entry || entry->deleted)
        {
          continue;
        }
      if (position > (target_position > (size_t)before ? target_position - (size_t)before : 0))
        {
          tiny_json_raw(json, ",");
        }
      tiny_json_init(&item);
      tiny_mcp_tool_append_message(&item, &entry->message);
      tiny_json_raw_n(json, item.data, item.length);
      tiny_json_free(&item);
    }
  tiny_json_raw(json, "],\"after\":[");
  for (position = target_position + 1;
       position < entry_count && position < target_position + 1 + (size_t)after;
       ++position)
    {
      const tiny_store_entry_t *entry = tiny_store_entry(context->store, entries[position]);
      tiny_json_buffer_t item;
      if (!entry || entry->deleted)
        {
          continue;
        }
      if (position > target_position + 1)
        {
          tiny_json_raw(json, ",");
        }
      tiny_json_init(&item);
      tiny_mcp_tool_append_message(&item, &entry->message);
      tiny_json_raw_n(json, item.data, item.length);
      tiny_json_free(&item);
    }
  tiny_json_raw(json, "]}");
}

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 *filters_json = tiny_mcp_tool_get_string(tool, "filters_json", 4096);
  int32_t requested_limit = 20;
  int32_t minimum_importance = 0;
  int32_t context_size = 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);
  mcpc_errcode_t context_status
      = mcpc_tool_get_tpropval_i32(tool, (const char8_t *)"context_size", &context_size);
  search_filters_t filters = { { false } };
  size_t limit, index, count = 0, candidate_count = 0, *candidates = NULL;
  result_item_t *items = NULL;
  tiny_json_buffer_t json;
  bool filters_ok;
  bool valid = context && query && query[0]
               && (limit_status == MCPC_EC_0 || limit_status == MCPC_EC_ANYDATA_NOTFOUND)
               && requested_limit >= 1 && requested_limit <= SEARCH_MAX_LIMIT;
  valid = valid && minimum_importance >= 0 && minimum_importance <= 100
          && (importance_status == MCPC_EC_0 || importance_status == MCPC_EC_ANYDATA_NOTFOUND)
          && context_size >= 0 && context_size <= SEARCH_MAX_CONTEXT_SIZE
          && (context_status == MCPC_EC_0 || context_status == MCPC_EC_ANYDATA_NOTFOUND);
  filters_ok = parse_filters_json(filters_json, &filters);
  if (!valid || !filters_ok)
    {
      tiny_mcp_tool_add_error(
          result, "VALIDATION_ERROR", "query, limit, filters or context_size is invalid");
      search_filters_free(&filters);
    }
  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]);
                  tiny_search_breakdown_t breakdown;
                  if (!entry || entry->deleted
                      || (session && session[0] && strcmp(entry->message.session_id, session) != 0))
                    {
                      continue;
                    }
                  if (entry->message.metadata.importance < minimum_importance)
                    {
                      continue;
                    }
                  if (!filters_match(&filters, &entry->message))
                    {
                      continue;
                    }
                  if (!tiny_search_score_breakdown(query, &entry->message, &breakdown)
                      || breakdown.total <= 0)
                    {
                      continue;
                    }
                  items[count].entry = entry;
                  items[count].score = breakdown.total;
                  items[count].content_score = breakdown.content;
                  items[count].summary_score = breakdown.summary;
                  items[count].tag_score = breakdown.tags;
                  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)
                {
                  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, "},\"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);
                  if (context_size > 0)
                    {
                      append_context_preview(
                          &item, context, &items[index].entry->message, context_size);
                    }
                  tiny_json_raw(&item, "}");
                  if (item.failed || json.length + item.length + 256 > SEARCH_RESPONSE_BUDGET)
                    {
                      tiny_json_free(&item);
                      if (item.failed)
                        {
                          json.failed = true;
                        }
                      break;
                    }
                  if (index)
                    {
                      tiny_json_raw(&json, ",");
                    }
                  tiny_json_raw_n(&json, item.data, item.length);
                  tiny_json_free(&item);
                }
              tiny_json_raw(&json, "]}");
              if (json.failed)
                {
                  tiny_mcp_tool_add_error(
                      result, "INTERNAL_ERROR", "search response allocation failed");
                }
              else
                {
                  tiny_mcp_tool_add_json(result, &json);
                }
              tiny_json_free(&json);
            }
        }
    }
  free(items);
  free(candidates);
  free(query);
  free(session);
  free(filters_json);
  search_filters_free(&filters);
}

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,
             "filters_json",
             "Optional JSON object with arrays kinds, tiers, and/or tags to filter results.",
             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,
             "context_size",
             "Messages before and after each result, from 0 to 3; 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;
}