// Written by retoor@molodetz.nl

// This code manages a collection of messages using JSON objects. It provides
// functions to retrieve all messages as a JSON array, add a new message with a
// specified role and content, and free the allocated resources.

// Uses the external library <json-c/json.h> for JSON manipulation

// MIT License
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions: The above copyright
// notice and this permission notice shall be included in all copies or
// substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS",
// WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
// TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE
// FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
// TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR
// THE USE OR OTHER DEALINGS IN THE SOFTWARE.

#ifndef R_MESSAGES_H
#define R_MESSAGES_H

#include "db_utils.h"
#include "json-c/json.h"
#include "tools.h"
#include "utils.h"
#include <string.h>
#include <unistd.h>

struct json_object *message_array = NULL;
static char session_id[256] = {0};
static char override_session_id[256] = {0};

struct json_object *message_list();

bool set_session_id(const char *name) {
  if (name == NULL || *name == '\0') {
    return false;
  }
  size_t len = strlen(name);
  if (len > 200) {
    return false;
  }
  for (const char *p = name; *p; p++) {
    if (*p == '/' || *p == '\\') {
      return false;
    }
  }
  strncpy(override_session_id, name, sizeof(override_session_id) - 1);
  override_session_id[sizeof(override_session_id) - 1] = '\0';
  return true;
}

const char *get_session_id() {
  if (session_id[0] == '\0') {
    if (override_session_id[0] != '\0') {
      strncpy(session_id, override_session_id, sizeof(session_id) - 1);
      session_id[sizeof(session_id) - 1] = '\0';
    }
  }
  return session_id;
}

void init_session_id() {
  if (session_id[0] != '\0') {
    return;
  }
  if (override_session_id[0] != '\0') {
    strncpy(session_id, override_session_id, sizeof(session_id) - 1);
    session_id[sizeof(session_id) - 1] = '\0';
    return;
  }
  char *tty = ttyname(STDIN_FILENO);
  if (tty) {
    unsigned long h = 5381;
    for (char *p = tty; *p; p++) {
      h = ((h << 5) + h) + (unsigned char)*p;
    }
    snprintf(session_id, sizeof(session_id), "%lu", h);
  } else {
    snprintf(session_id, sizeof(session_id), "%d", getppid());
  }
}

char *get_session_db_key() {
  static char key[512];
  init_session_id();
  snprintf(key, sizeof(key), "session:%s", session_id);
  return key;
}

bool messages_save() {
  char *key = get_session_db_key();
  const char *json_str = json_object_to_json_string_ext(message_list(), JSON_C_TO_STRING_PLAIN);
  if (!json_str) {
    return false;
  }
  json_object *result = db_set(key, json_str);
  if (result) {
    json_object_put(result);
    return true;
  }
  return false;
}

bool messages_load_conversation() {
  char *key = get_session_db_key();
  json_object *result = db_get(key);
  if (!result) {
    return false;
  }
  json_object *value_obj;
  if (!json_object_object_get_ex(result, "value", &value_obj)) {
    json_object_put(result);
    return false;
  }
  const char *json_str = json_object_get_string(value_obj);
  if (!json_str) {
    json_object_put(result);
    return false;
  }
  struct json_object *loaded = json_tokener_parse(json_str);
  json_object_put(result);
  if (!loaded || !json_object_is_type(loaded, json_type_array)) {
    if (loaded) {
      json_object_put(loaded);
    }
    return false;
  }
  int len = json_object_array_length(loaded);
  int added = 0;
  for (int i = 0; i < len; i++) {
    struct json_object *msg = json_object_array_get_idx(loaded, i);
    struct json_object *role_obj;
    if (json_object_object_get_ex(msg, "role", &role_obj)) {
      const char *role = json_object_get_string(role_obj);
      if (role && strcmp(role, "system") != 0) {
        json_object_array_add(message_list(), json_object_get(msg));
        added++;
      }
    } else {
      json_object_array_add(message_list(), json_object_get(msg));
      added++;
    }
  }
  json_object_put(loaded);
  return added > 0;
}

struct json_object *message_list() {
  if (!message_array) {
    message_array = json_object_new_array();
  }
  return message_array;
}
bool messages_remove_last() {
  struct json_object *messages = message_list();
  int size = json_object_array_length(messages);
  if (size) {
    json_object_array_del_idx(messages, size - 1, 1);
    messages_save();
    return true;
  }
  return false;
}

void messages_remove() {
  if (message_array) {
    json_object_put(message_array);
    message_array = NULL;
  }
  message_list();
  messages_save();
}

struct json_object *message_add_tool_call(struct json_object *message) {
  struct json_object *messages = message_list();
  json_object_array_add(messages, message);
  messages_save();
  return message;
}

struct json_object *message_add_tool_result(const char *tool_call_id,
                                            char *tool_result) {
  struct json_object *messages = message_list();
  struct json_object *message = json_object_new_object();
  if (message == NULL) {
    return NULL;
  }

  json_object_object_add(message, "tool_call_id",
                         json_object_new_string(tool_call_id));

  size_t result_len = strlen(tool_result);
  if (result_len > 104000) {
    result_len = 104000;
  }
  json_object_object_add(message, "tool_result",
                         json_object_new_string_len(tool_result, (int)result_len));

  json_object_array_add(messages, message);
  messages_save();
  return message;
}

void message_add_object(json_object *message) {
  struct json_object *messages = message_list();
  json_object_array_add(messages, message);
  messages_save();
}

struct json_object *message_add(const char *role, const char *content);

struct json_object *message_add(const char *role, const char *content) {
  struct json_object *messages = message_list();
  struct json_object *message = json_object_new_object();
  if (message == NULL) {
    return NULL;
  }
  json_object_object_add(message, "role", json_object_new_string(role));

  if (content) {
    size_t content_len = strlen(content);
    if (content_len > 1048570) {
      content_len = 1048570;
    }
    json_object_object_add(message, "content",
                           json_object_new_string_len(content, (int)content_len));
  }
  if (!strcmp(role, "user")) {
    json_object_object_add(message, "tools", tools_descriptions());
    json_object_object_add(message, "parallel_tool_calls",
                           json_object_new_boolean(true));
  }

  json_object_array_add(messages, message);
  if (strcmp(role, "system") != 0) {
    messages_save();
  }
  return message;
}

char *message_json() {
  return (char *)json_object_to_json_string_ext(message_list(),
                                                JSON_C_TO_STRING_PRETTY);
}

void message_free() {
  if (message_array) {
    json_object_put(message_array);
    message_array = NULL;
  }
}

#endif