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

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

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

#include "mjson.h"
#include "storage/store.h"

#define IMPORT_MAX_LINE (TINY_CONTENT_MAX + 4096U)

typedef struct import_record_list
{
  tiny_import_record_t *items;
  size_t count;
  size_t capacity;
} import_record_list_t;

static void
import_record_list_free(import_record_list_t *list)
{
  size_t index;
  if (!list)
    {
      return;
    }
  for (index = 0; index < list->count; ++index)
    {
      tiny_message_free(&list->items[index].message);
    }
  free(list->items);
  list->items = NULL;
  list->count = 0;
  list->capacity = 0;
}

static bool
import_record_list_push(import_record_list_t *list)
{
  tiny_import_record_t *resized;
  size_t capacity;
  if (list->count < list->capacity)
    {
      return true;
    }
  capacity = list->capacity == 0 ? 32 : list->capacity * 2;
  resized = realloc(list->items, capacity * sizeof(*resized));
  if (!resized)
    {
      return false;
    }
  list->items = resized;
  list->capacity = capacity;
  return true;
}

static bool
import_get_string(const char *json, int len, const char *path, char *out, size_t out_size)
{
  int decoded = mjson_get_string(json, len, path, out, (int)out_size);
  return decoded > 0 && (size_t)decoded < out_size;
}

static bool
import_get_i64(const char *json, int len, const char *path, int64_t *value)
{
  double number;
  if (mjson_get_number(json, len, path, &number) != 1)
    {
      return false;
    }
  *value = (int64_t)number;
  return true;
}

static bool
import_get_u64(const char *json, int len, const char *path, uint64_t *value)
{
  double number;
  if (mjson_get_number(json, len, path, &number) != 1 || number < 0)
    {
      return false;
    }
  *value = (uint64_t)number;
  return true;
}

static bool
import_parse_string_array(const char *json,
                          int len,
                          const char *path,
                          char ***output,
                          size_t *count,
                          size_t max_count,
                          size_t max_length)
{
  int offset = 0, key_offset, key_length, value_offset, value_length, value_type;
  char *array_json = NULL;
  int array_len = 0;
  size_t index = 0;
  *output = NULL;
  *count = 0;
  if (mjson_find(json, len, path, &array_json, &array_len) == 0)
    {
      return true;
    }
  if (array_json == NULL || mjson_find(array_json, array_len, "$", NULL, NULL) != MJSON_TOK_ARRAY)
    {
      return false;
    }
  *output = calloc(max_count, sizeof(char *));
  if (!*output)
    {
      return false;
    }
  while ((offset = mjson_next(array_json,
                              array_len,
                              offset,
                              &key_offset,
                              &key_length,
                              &value_offset,
                              &value_length,
                              &value_type))
         != 0)
    {
      char value[512];
      int decoded;
      if (value_type != MJSON_TOK_STRING || index >= max_count
          || (decoded = mjson_get_string(
                  array_json + value_offset, value_length, "$", value, sizeof(value)))
                 <= 0
          || (size_t)decoded > max_length || decoded >= (int)sizeof(value))
        {
          break;
        }
      (*output)[index] = malloc((size_t)decoded + 1);
      if (!(*output)[index])
        {
          break;
        }
      memcpy((*output)[index], value, (size_t)decoded);
      (*output)[index][decoded] = '\0';
      index++;
    }
  *count = index;
  if (index > 0 && offset != 0)
    {
      /* mjson_next returned non-zero because of a parse/validation failure. */
      return false;
    }
  return true;
}

static bool
import_parse_relations(const char *json, int len, const char *path, tiny_metadata_t *metadata)
{
  int offset = 0, key_offset, key_length, value_offset, value_length, value_type;
  char *array_json = NULL;
  int array_len = 0;
  size_t index = 0;
  if (mjson_find(json, len, path, &array_json, &array_len) == 0)
    {
      return true;
    }
  if (array_json == NULL || mjson_find(array_json, array_len, "$", NULL, NULL) != MJSON_TOK_ARRAY)
    {
      return false;
    }
  metadata->relations = calloc(TINY_RELATION_COUNT_MAX, sizeof(tiny_relation_t));
  if (!metadata->relations)
    {
      return false;
    }
  while ((offset = mjson_next(array_json,
                              array_len,
                              offset,
                              &key_offset,
                              &key_length,
                              &value_offset,
                              &value_length,
                              &value_type))
         != 0)
    {
      char type_name[32];
      char message_id[32];
      int decoded;
      tiny_relation_t *relation;
      if (value_type != MJSON_TOK_OBJECT || index >= TINY_RELATION_COUNT_MAX
          || mjson_get_string(
                 array_json + value_offset, value_length, "$.type", type_name, sizeof(type_name))
                 <= 0
          || (decoded = mjson_get_string(array_json + value_offset,
                                         value_length,
                                         "$.message_id",
                                         message_id,
                                         sizeof(message_id)))
                 <= 0)
        {
          return false;
        }
      relation = &metadata->relations[index];
      if (!tiny_relation_type_parse(type_name, &relation->type)
          || !tiny_message_id_valid(message_id))
        {
          return false;
        }
      memcpy(relation->message_id, message_id, (size_t)decoded);
      relation->message_id[decoded] = '\0';
      index++;
    }
  metadata->relation_count = index;
  return true;
}

static bool
import_parse_metadata(const char *json, int len, tiny_message_t *message)
{
  char tier_name[32] = { 0 };
  char kind_name[32] = { 0 };
  char tool_call_id[TINY_TOOL_CALL_ID_MAX + 1] = { 0 };
  char summary[TINY_SUMMARY_MAX + 1] = { 0 };
  double importance_number = 0;
  bool has_tool_call_id, has_summary, has_importance;
  has_tool_call_id = mjson_find(json, len, "$.metadata.tool_call_id", NULL, NULL) != 0
                     && import_get_string(
                         json, len, "$.metadata.tool_call_id", tool_call_id, sizeof(tool_call_id));
  has_summary = mjson_find(json, len, "$.metadata.summary", NULL, NULL) != 0
                && import_get_string(json, len, "$.metadata.summary", summary, sizeof(summary));
  has_importance = mjson_get_number(json, len, "$.metadata.importance", &importance_number) == 1;
  if (!import_get_string(json, len, "$.metadata.memory_tier", tier_name, sizeof(tier_name))
      || !import_get_string(json, len, "$.metadata.memory_kind", kind_name, sizeof(kind_name))
      || !import_parse_string_array(json,
                                    len,
                                    "$.metadata.tags",
                                    &message->metadata.tags,
                                    &message->metadata.tag_count,
                                    TINY_TAG_COUNT_MAX,
                                    TINY_TAG_MAX)
      || !import_parse_relations(json, len, "$.metadata.relations", &message->metadata)
      || !import_parse_string_array(json,
                                    len,
                                    "$.metadata.source_message_ids",
                                    &message->metadata.source_message_ids,
                                    &message->metadata.source_count,
                                    TINY_SOURCE_COUNT_MAX,
                                    TINY_MESSAGE_ID_LENGTH))
    {
      return false;
    }
  if (!tiny_tier_parse(tier_name, &message->metadata.memory_tier)
      || !tiny_kind_parse(kind_name, &message->metadata.memory_kind))
    {
      return false;
    }
  if (has_tool_call_id)
    {
      message->metadata.tool_call_id = malloc(strlen(tool_call_id) + 1);
      if (!message->metadata.tool_call_id)
        {
          return false;
        }
      memcpy(message->metadata.tool_call_id, tool_call_id, strlen(tool_call_id) + 1);
    }
  if (has_summary)
    {
      message->metadata.summary = malloc(strlen(summary) + 1);
      if (!message->metadata.summary)
        {
          return false;
        }
      memcpy(message->metadata.summary, summary, strlen(summary) + 1);
    }
  if (has_importance)
    {
      if (importance_number < 0 || importance_number > 100)
        {
          return false;
        }
      message->metadata.importance = (uint8_t)importance_number;
    }
  return true;
}

static bool
import_manifest_valid(const char *line, int len)
{
  char type_name[64];
  double version = 0;
  if (!import_get_string(line, len, "$.type", type_name, sizeof(type_name))
      || strcmp(type_name, "tiny-convmem-manifest") != 0
      || mjson_get_number(line, len, "$.format_version", &version) != 1)
    {
      return false;
    }
  return (int)version == 2;
}

static bool
import_parse_del_record(const char *line, int len, tiny_import_record_t *record)
{
  char type_name[32];
  tiny_message_init(&record->message);
  if (!import_get_string(line, len, "$.type", type_name, sizeof(type_name))
      || strcmp(type_name, "DEL") != 0
      || !import_get_string(
          line, len, "$.message_id", record->message.message_id, sizeof(record->message.message_id))
      || !tiny_message_id_valid(record->message.message_id))
    {
      return false;
    }
  record->deleted = true;
  return true;
}

static bool
import_parse_message_record(const char *line, int len, tiny_import_record_t *record)
{
  char role_name[32];
  char session_id[TINY_SESSION_ID_MAX + 1];
  char content_buffer[4096];
  tiny_message_init(&record->message);
  record->deleted = false;
  if (!import_get_string(
          line, len, "$.message_id", record->message.message_id, sizeof(record->message.message_id))
      || !import_get_string(line, len, "$.session_id", session_id, sizeof(session_id))
      || !import_get_string(line, len, "$.role", role_name, sizeof(role_name))
      || !import_get_i64(line, len, "$.timestamp", &record->message.timestamp_ms)
      || !import_get_i64(line, len, "$.updated_at", &record->message.updated_at_ms)
      || !import_get_u64(line, len, "$.revision", &record->message.revision))
    {
      return false;
    }
  if (!tiny_message_id_valid(record->message.message_id)
      || !tiny_role_parse(role_name, &record->message.role) || record->message.timestamp_ms < 0
      || record->message.updated_at_ms < 0 || record->message.revision == 0
      || record->message.revision > INT32_MAX)
    {
      return false;
    }
  if (mjson_find(line, len, "$.content", NULL, NULL) == 0)
    {
      return false;
    }
  if (mjson_get_string(line, len, "$.content", content_buffer, sizeof(content_buffer)) > 0)
    {
      size_t length = strlen(content_buffer);
      if (length == 0 || length > TINY_CONTENT_MAX)
        {
          return false;
        }
      record->message.content = malloc(length + 1);
      if (!record->message.content)
        {
          return false;
        }
      memcpy(record->message.content, content_buffer, length + 1);
    }
  else
    {
      const char *content = NULL;
      int content_len = 0;
      if (mjson_find(line, len, "$.content", &content, &content_len) == 0 || !content
          || content_len < 2)
        {
          return false;
        }
      if ((size_t)(content_len - 2) == 0 || (size_t)(content_len - 2) > TINY_CONTENT_MAX)
        {
          return false;
        }
      record->message.content = malloc((size_t)content_len - 1);
      if (!record->message.content)
        {
          return false;
        }
      memcpy(record->message.content, content + 1, (size_t)content_len - 2);
      record->message.content[(size_t)content_len - 2] = '\0';
    }
  if (!tiny_utf8_valid(record->message.content, strlen(record->message.content)))
    {
      return false;
    }
  record->message.session_id = malloc(strlen(session_id) + 1);
  if (!record->message.session_id)
    {
      return false;
    }
  memcpy(record->message.session_id, session_id, strlen(session_id) + 1);
  if (!import_parse_metadata(line, len, &record->message)
      || !tiny_metadata_valid(&record->message.metadata))
    {
      tiny_message_free(&record->message);
      tiny_message_init(&record->message);
      return false;
    }
  return true;
}

static bool
import_parse_jsonl(const char *jsonl, import_record_list_t *records)
{
  const char *line_start = jsonl;
  const char *line_end;
  bool manifest_seen = false;
  while (*line_start)
    {
      int length;
      tiny_import_record_t record;
      line_end = strchr(line_start, '\n');
      if (line_end == NULL)
        {
          line_end = line_start + strlen(line_start);
        }
      length = (int)(line_end - line_start);
      if (length > 0)
        {
          if (!manifest_seen)
            {
              if (!import_manifest_valid(line_start, length))
                {
                  return false;
                }
              manifest_seen = true;
            }
          else if (import_parse_del_record(line_start, length, &record))
            {
              if (!import_record_list_push(records))
                {
                  return false;
                }
              records->items[records->count++] = record;
            }
          else if (import_parse_message_record(line_start, length, &record))
            {
              if (!import_record_list_push(records))
                {
                  tiny_message_free(&record.message);
                  return false;
                }
              records->items[records->count++] = record;
            }
          else
            {
              return false;
            }
        }
      if (*line_end == '\0')
        {
          break;
        }
      line_start = line_end + 1;
    }
  return manifest_seen;
}

static void
call_import_messages(const mcpc_tool_t *tool, mcpc_ucbr_t *result)
{
  tiny_mcp_tools_context_t *context = tiny_mcp_tools_context_current();
  char *jsonl = tiny_mcp_tool_get_string(tool, "jsonl", TINY_CONTENT_MAX);
  char *policy_name = tiny_mcp_tool_get_string(tool, "policy", 16);
  tiny_import_policy_t policy = TINY_IMPORT_OVERWRITE;
  import_record_list_t records = { NULL, 0, 0 };
  size_t imported = 0, rejected = 0;
  tiny_json_buffer_t json;
  tiny_store_status_t status;

  if (!context || !jsonl)
    {
      tiny_mcp_tool_add_error(result, "VALIDATION_ERROR", "invalid import_messages arguments");
      free(jsonl);
      free(policy_name);
      return;
    }
  if (policy_name && policy_name[0])
    {
      if (strcmp(policy_name, "overwrite") == 0)
        {
          policy = TINY_IMPORT_OVERWRITE;
        }
      else if (strcmp(policy_name, "reject") == 0)
        {
          policy = TINY_IMPORT_REJECT;
        }
      else
        {
          tiny_mcp_tool_add_error(result, "VALIDATION_ERROR", "invalid import policy");
          free(jsonl);
          free(policy_name);
          return;
        }
    }
  if (!import_parse_jsonl(jsonl, &records))
    {
      tiny_mcp_tool_add_error(result, "VALIDATION_ERROR", "invalid import JSONL");
      import_record_list_free(&records);
      free(jsonl);
      free(policy_name);
      return;
    }
  status = tiny_store_import(
      context->store, records.items, records.count, policy, &imported, &rejected);
  import_record_list_free(&records);
  tiny_json_init(&json);
  if (status == TINY_STORE_OK)
    {
      tiny_json_raw(&json, "{\"status\":\"OK\",\"imported\":");
      tiny_json_u64(&json, (uint64_t)imported);
      tiny_json_raw(&json, ",\"rejected\":");
      tiny_json_u64(&json, (uint64_t)rejected);
      tiny_json_raw(&json, "}");
    }
  else if (status == TINY_STORE_CONFLICT)
    {
      tiny_json_raw(&json,
                    "{\"status\":\"ERROR\",\"error\":{\"code\":\"CONFLICT\","
                    "\"message\":\"import conflict with existing records\"}}");
    }
  else if (status == TINY_STORE_CORRUPT)
    {
      tiny_json_raw(&json,
                    "{\"status\":\"ERROR\",\"error\":{\"code\":\"CORRUPT_STORE\","
                    "\"message\":\"import JSONL contained invalid records\"}}");
    }
  else
    {
      tiny_json_raw(&json,
                    "{\"status\":\"ERROR\",\"error\":{\"code\":\"INTERNAL_ERROR\","
                    "\"message\":\"import failed\"}}");
    }
  tiny_mcp_tool_add_json(result, &json);
  tiny_json_free(&json);
  free(jsonl);
  free(policy_name);
}

bool
tiny_mcp_register_import_messages(mcpc_server_t *server)
{
  mcpc_tool_t *tool = tiny_mcp_tool_new(
      "import_messages",
      "Atomically import a JSONL export from read_messages view=export. The manifest line is "
      "required; subsequent lines are PUT messages or DEL tombstones. Conflicting message IDs "
      "are resolved by the policy argument.",
      call_import_messages);
  return tool
         && tiny_mcp_tool_add_property(tool,
                                       "jsonl",
                                       "JSONL string containing a manifest line followed by PUT "
                                       "message lines and optional DEL tombstone lines.",
                                       MCPC_U8STR,
                                       true)
         && tiny_mcp_tool_add_property(tool,
                                       "policy",
                                       "Conflict policy: overwrite (default) replaces existing "
                                       "message IDs, reject refuses duplicate IDs.",
                                       MCPC_U8STR,
                                       false)
         && mcpc_server_add_tool(server, tool) == MCPC_EC_0;
}