#include "cli/cli.h"

#include "domain/message.h"
#include "mcp/json.h"
#include "storage/store.h"

#include "mjson.h"

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

#define IMPORT_MAX_LINE TINY_CLI_LINE_MAX
#define IMPORT_MAX_RECORDS 100000U
#define IMPORT_MAX_DECODED_BYTES (UINT64_C(4) * 1024U * 1024U * 1024U)

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 int
import_hex_digit(char value)
{
  if (value >= '0' && value <= '9')
    {
      return value - '0';
    }
  if (value >= 'a' && value <= 'f')
    {
      return value - 'a' + 10;
    }
  if (value >= 'A' && value <= 'F')
    {
      return value - 'A' + 10;
    }
  return -1;
}

static bool
import_decode_json_string(const char *source, size_t length, char *output, size_t *output_length)
{
  size_t input = 0, used = 0;
  while (input < length)
    {
      uint32_t codepoint;
      unsigned count;
      if (source[input] != '\\')
        {
          output[used++] = source[input++];
          continue;
        }
      if (++input >= length)
        {
          return false;
        }
      switch (source[input++])
        {
        case '"':
          output[used++] = '"';
          continue;
        case '\\':
          output[used++] = '\\';
          continue;
        case '/':
          output[used++] = '/';
          continue;
        case 'b':
          output[used++] = '\b';
          continue;
        case 'f':
          output[used++] = '\f';
          continue;
        case 'n':
          output[used++] = '\n';
          continue;
        case 'r':
          output[used++] = '\r';
          continue;
        case 't':
          output[used++] = '\t';
          continue;
        case 'u':
          break;
        default:
          return false;
        }
      if (input + 4 > length)
        {
          return false;
        }
      codepoint = 0;
      for (count = 0; count < 4; ++count)
        {
          int digit = import_hex_digit(source[input++]);
          if (digit < 0)
            {
              return false;
            }
          codepoint = (codepoint << 4) | (uint32_t)digit;
        }
      if (codepoint >= 0xd800 && codepoint <= 0xdbff)
        {
          uint32_t low = 0;
          if (input + 6 > length || source[input] != '\\' || source[input + 1] != 'u')
            {
              return false;
            }
          input += 2;
          for (count = 0; count < 4; ++count)
            {
              int digit = import_hex_digit(source[input++]);
              if (digit < 0)
                {
                  return false;
                }
              low = (low << 4) | (uint32_t)digit;
            }
          if (low < 0xdc00 || low > 0xdfff)
            {
              return false;
            }
          codepoint = UINT32_C(0x10000) + ((codepoint - 0xd800U) << 10) + low - 0xdc00U;
        }
      else if (codepoint >= 0xdc00 && codepoint <= 0xdfff)
        {
          return false;
        }
      if (codepoint <= 0x7f)
        {
          output[used++] = (char)codepoint;
        }
      else if (codepoint <= 0x7ff)
        {
          output[used++] = (char)(0xc0U | (codepoint >> 6));
          output[used++] = (char)(0x80U | (codepoint & 0x3fU));
        }
      else if (codepoint <= 0xffff)
        {
          output[used++] = (char)(0xe0U | (codepoint >> 12));
          output[used++] = (char)(0x80U | ((codepoint >> 6) & 0x3fU));
          output[used++] = (char)(0x80U | (codepoint & 0x3fU));
        }
      else
        {
          output[used++] = (char)(0xf0U | (codepoint >> 18));
          output[used++] = (char)(0x80U | ((codepoint >> 12) & 0x3fU));
          output[used++] = (char)(0x80U | ((codepoint >> 6) & 0x3fU));
          output[used++] = (char)(0x80U | (codepoint & 0x3fU));
        }
    }
  output[used] = '\0';
  *output_length = used;
  return true;
}

static bool
import_get_allocated_string(const char *json,
                            int len,
                            const char *path,
                            size_t max_length,
                            bool allow_empty,
                            char **out,
                            size_t *decoded_length)
{
  const char *token = NULL;
  int token_length = 0;
  size_t decoded;
  char *value;
  *out = NULL;
  if (decoded_length)
    {
      *decoded_length = 0;
    }
  if (mjson_find(json, len, path, &token, &token_length) != MJSON_TOK_STRING || token_length < 2
      || token_length == INT_MAX)
    {
      return false;
    }
  value = malloc((size_t)token_length);
  if (!value)
    {
      return false;
    }
  if (!import_decode_json_string(token + 1, (size_t)token_length - 2U, value, &decoded)
      || (!allow_empty && decoded == 0) || decoded > max_length || strlen(value) != decoded
      || !tiny_utf8_valid(value, decoded))
    {
      free(value);
      return false;
    }
  *out = value;
  if (decoded_length)
    {
      *decoded_length = decoded;
    }
  return true;
}

static bool
import_get_integer_token(const char *json,
                         int len,
                         const char *path,
                         bool allow_negative,
                         uint64_t maximum,
                         uint64_t *magnitude,
                         bool *negative)
{
  const char *token = NULL;
  int token_length = 0;
  int index = 0;
  uint64_t value = 0;
  if (mjson_find(json, len, path, &token, &token_length) != MJSON_TOK_NUMBER || token_length <= 0)
    {
      return false;
    }
  *negative = token[0] == '-';
  if (*negative)
    {
      if (!allow_negative || ++index == token_length)
        {
          return false;
        }
    }
  if (token[index] == '0' && index + 1 != token_length)
    {
      return false;
    }
  for (; index < token_length; ++index)
    {
      unsigned digit;
      if (token[index] < '0' || token[index] > '9')
        {
          return false;
        }
      digit = (unsigned)(token[index] - '0');
      if (value > (maximum - digit) / 10U)
        {
          return false;
        }
      value = value * 10U + digit;
    }
  *magnitude = value;
  return true;
}

static bool
import_get_i64(const char *json, int len, const char *path, int64_t *value)
{
  uint64_t magnitude;
  bool negative;
  if (!import_get_integer_token(
          json, len, path, true, (uint64_t)INT64_MAX + 1U, &magnitude, &negative))
    {
      return false;
    }
  if (negative)
    {
      *value = magnitude == (uint64_t)INT64_MAX + 1U ? INT64_MIN : -(int64_t)magnitude;
    }
  else
    {
      if (magnitude > INT64_MAX)
        {
          return false;
        }
      *value = (int64_t)magnitude;
    }
  return true;
}

static bool
import_get_u64(const char *json, int len, const char *path, uint64_t *value)
{
  bool negative;
  if (!import_get_integer_token(json, len, path, false, UINT64_MAX, value, &negative))
    {
      return false;
    }
  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 = NULL;
      size_t decoded = 0;
      if (value_type != MJSON_TOK_STRING || index >= max_count
          || !import_get_allocated_string(
              array_json + value_offset, value_length, "$", max_length, false, &value, &decoded))
        {
          break;
        }
      (*output)[index] = value;
      index++;
    }
  *count = index;
  if (index > 0 && offset != 0)
    {
      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 = NULL;
      char *message_id = NULL;
      size_t decoded = 0;
      tiny_relation_t *relation;
      if (value_type != MJSON_TOK_OBJECT || index >= TINY_RELATION_COUNT_MAX
          || !import_get_allocated_string(
              array_json + value_offset, value_length, "$.type", 31, false, &type_name, NULL)
          || !import_get_allocated_string(array_json + value_offset,
                                          value_length,
                                          "$.message_id",
                                          TINY_MESSAGE_ID_LENGTH,
                                          false,
                                          &message_id,
                                          &decoded))
        {
          free(type_name);
          free(message_id);
          return false;
        }
      relation = &metadata->relations[index];
      if (!tiny_relation_type_parse(type_name, &relation->type)
          || !tiny_message_id_valid(message_id))
        {
          free(type_name);
          free(message_id);
          return false;
        }
      memcpy(relation->message_id, message_id, (size_t)decoded);
      relation->message_id[decoded] = '\0';
      free(type_name);
      free(message_id);
      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 = NULL;
  char *summary = NULL;
  uint64_t importance_number = 0;
  bool negative = false;
  bool has_tool_call_id, has_summary, has_importance;
  has_tool_call_id = mjson_find(json, len, "$.metadata.tool_call_id", NULL, NULL) != 0;
  has_summary = mjson_find(json, len, "$.metadata.summary", NULL, NULL) != 0;
  has_importance = mjson_find(json, len, "$.metadata.importance", NULL, NULL) != 0;
  if ((has_tool_call_id
       && !import_get_allocated_string(
           json, len, "$.metadata.tool_call_id", TINY_TOOL_CALL_ID_MAX, false, &tool_call_id, NULL))
      || (has_summary
          && !import_get_allocated_string(
              json, len, "$.metadata.summary", TINY_SUMMARY_MAX, false, &summary, NULL))
      || (has_importance
          && !import_get_integer_token(
              json, len, "$.metadata.importance", false, 100, &importance_number, &negative)))
    {
      free(tool_call_id);
      free(summary);
      return false;
    }
  message->metadata.tool_call_id = tool_call_id;
  message->metadata.summary = summary;
  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_importance)
    {
      message->metadata.importance = (uint8_t)importance_number;
    }
  return true;
}

static bool
import_manifest_valid(const char *line, int len)
{
  char type_name[64];
  uint64_t version = 0;
  if (!import_get_string(line, len, "$.type", type_name, sizeof(type_name))
      || strcmp(type_name, "tiny-convmem-manifest") != 0
      || !import_get_u64(line, len, "$.format_version", &version))
    {
      return false;
    }
  return 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 = NULL;
  char *content = NULL;
  size_t content_length = 0;
  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, "$.role", role_name, sizeof(role_name))
      || !import_get_allocated_string(
          line, len, "$.session_id", TINY_SESSION_ID_MAX, false, &session_id, NULL)
      || !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))
    {
      free(session_id);
      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)
    {
      free(session_id);
      return false;
    }
  if (!import_get_allocated_string(
          line, len, "$.content", TINY_CONTENT_MAX, false, &content, &content_length))
    {
      free(session_id);
      return false;
    }
  (void)content_length;
  record->message.content = content;
  record->message.session_id = session_id;
  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 uint64_t
import_record_decoded_bytes(const tiny_import_record_t *record)
{
  uint64_t total = 0;
  size_t index;
  const tiny_message_t *message = &record->message;
  if (record->deleted)
    {
      return TINY_MESSAGE_ID_LENGTH;
    }
  total = strlen(message->session_id) + strlen(message->content) + TINY_MESSAGE_ID_LENGTH;
  total += message->metadata.tool_call_id ? strlen(message->metadata.tool_call_id) : 0;
  total += message->metadata.summary ? strlen(message->metadata.summary) : 0;
  for (index = 0; index < message->metadata.tag_count; ++index)
    {
      total += strlen(message->metadata.tags[index]);
    }
  total += message->metadata.relation_count * TINY_MESSAGE_ID_LENGTH;
  total += message->metadata.source_count * TINY_MESSAGE_ID_LENGTH;
  return total;
}

static bool
import_parse_stream(FILE *file, import_record_list_t *records)
{
  char *line = malloc((size_t)IMPORT_MAX_LINE + 2U);
  bool manifest_seen = false;
  uint64_t decoded_total = 0;
  if (!line)
    {
      return false;
    }
  while (fgets(line, (int)IMPORT_MAX_LINE + 2, file) != NULL)
    {
      size_t length = strlen(line);
      tiny_import_record_t record;
      uint64_t decoded;
      if (length == 0 || length > IMPORT_MAX_LINE
          || (length == IMPORT_MAX_LINE + 1U && line[length - 1] != '\n')
          || (line[length - 1] != '\n' && !feof(file)))
        {
          free(line);
          return false;
        }
      if (length > 0 && line[length - 1] == '\n')
        {
          line[--length] = '\0';
        }
      if (length > 0 && line[length - 1] == '\r')
        {
          line[--length] = '\0';
        }
      if (length == 0)
        {
          continue;
        }
      if (!manifest_seen)
        {
          if (!import_manifest_valid(line, (int)length))
            {
              free(line);
              return false;
            }
          manifest_seen = true;
          continue;
        }
      if (records->count >= IMPORT_MAX_RECORDS)
        {
          free(line);
          return false;
        }
      if (!import_parse_del_record(line, (int)length, &record)
          && !import_parse_message_record(line, (int)length, &record))
        {
          free(line);
          return false;
        }
      decoded = import_record_decoded_bytes(&record);
      if (decoded > IMPORT_MAX_DECODED_BYTES - decoded_total || !import_record_list_push(records))
        {
          tiny_message_free(&record.message);
          free(line);
          return false;
        }
      decoded_total += decoded;
      records->items[records->count++] = record;
    }
  if (ferror(file))
    {
      free(line);
      return false;
    }
  free(line);
  return manifest_seen;
}

int
tiny_cli_import(tiny_store_t *store, const tiny_cli_options_t *options)
{
  import_record_list_t records = { NULL, 0, 0 };
  size_t imported = 0, rejected = 0;
  tiny_store_status_t status;
  tiny_import_policy_t policy = TINY_IMPORT_OVERWRITE;
  FILE *file = NULL;
  int result = 1;

  if (!options || !options->import_path)
    {
      fprintf(stderr, "import: path is required\n");
      return 1;
    }

  if (options->policy_name && options->policy_name[0])
    {
      if (strcmp(options->policy_name, "overwrite") == 0)
        {
          policy = TINY_IMPORT_OVERWRITE;
        }
      else if (strcmp(options->policy_name, "reject") == 0)
        {
          policy = TINY_IMPORT_REJECT;
        }
      else
        {
          fprintf(stderr, "import: invalid policy '%s'\n", options->policy_name);
          return 1;
        }
    }

  if (fopen_s(&file, options->import_path, "rb") != 0 || !file)
    {
      fprintf(stderr, "import: unable to open '%s'\n", options->import_path);
      return 1;
    }

  if (!import_parse_stream(file, &records))
    {
      fprintf(stderr, "import: invalid JSONL in '%s'\n", options->import_path);
      fclose(file);
      import_record_list_free(&records);
      return 1;
    }
  if (fclose(file) != 0)
    {
      fprintf(stderr, "import: unable to close '%s'\n", options->import_path);
      import_record_list_free(&records);
      return 1;
    }

  status = tiny_store_import(store, records.items, records.count, policy, &imported, &rejected);
  import_record_list_free(&records);

  if (status == TINY_STORE_OK)
    {
      printf("{\"imported\":%zu,\"rejected\":%zu}\n", imported, rejected);
      fflush(stdout);
      result = 0;
    }
  else if (status == TINY_STORE_CONFLICT)
    {
      fprintf(stderr, "import: conflict with existing records\n");
    }
  else if (status == TINY_STORE_CORRUPT)
    {
      fprintf(stderr, "import: invalid records\n");
    }
  else
    {
      fprintf(stderr, "import: failed\n");
    }

  return result;
}