#include "mrjunejune/conversation_store.h"

#include "deita/deita.h"

#include <fcntl.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

struct Conversation_Store {
  Deita_Connection *p_connection;
  pthread_mutex_t mutex;
};

boolean Conversation_Store_Generate_UUID(char output[37])
{
  uint8 bytes[16];
  int fd = open("/dev/urandom", O_RDONLY);
  if (fd < 0)
    return FALSE;
  size_t offset = 0;
  while (offset < sizeof(bytes))
  {
    ssize_t amount = read(fd, bytes + offset, sizeof(bytes) - offset);
    if (amount <= 0)
    {
      close(fd);
      return FALSE;
    }
    offset += (size_t)amount;
  }
  close(fd);

  bytes[6] = (uint8)((bytes[6] & 0x0f) | 0x40);
  bytes[8] = (uint8)((bytes[8] & 0x3f) | 0x80);
  snprintf(
      output,
      37,
      "%02x%02x%02x%02x-%02x%02x-%02x%02x-%02x%02x-"
      "%02x%02x%02x%02x%02x%02x",
      bytes[0], bytes[1], bytes[2], bytes[3],
      bytes[4], bytes[5], bytes[6], bytes[7],
      bytes[8], bytes[9], bytes[10], bytes[11],
      bytes[12], bytes[13], bytes[14], bytes[15]);
  return TRUE;
}

static char *Conversation_Store_Copy_Text(
    const char *value,
    Dowa_Arena *p_arena)
{
  const char *source = value ? value : "";
  size_t length = strlen(source);
  char *copy = Dowa_Arena_Allocate(p_arena, length + 1);
  if (!copy)
    return NULL;
  memcpy(copy, source, length + 1);
  return copy;
}

static boolean Conversation_Store_Exists_Locked(
    Conversation_Store *p_store,
    const char *conversation_id)
{
  Dowa_Arena *p_arena = Dowa_Arena_Create(1024);
  if (!p_arena)
    return FALSE;
  const char *parameters[] = {conversation_id};
  Deita_Result_Set *p_result = Deita_Query_Execute_Prepared(
      p_store->p_connection,
      "SELECT 1 FROM conversations WHERE id = ? AND status != 'deleted'",
      1,
      parameters,
      p_arena);
  boolean exists = p_result && Deita_Result_Set_Next(p_result);
  if (p_result)
    Deita_Result_Set_Free(p_result);
  Dowa_Arena_Free(p_arena);
  return exists;
}

static Conversation_Store_Result Conversation_Store_Rollback(
    Conversation_Store *p_store,
    Conversation_Store_Result result)
{
  Deita_Query_Execute_Update(p_store->p_connection, "ROLLBACK");
  return result;
}

Conversation_Store *Conversation_Store_Create(const char *database_path)
{
  if (!database_path)
    return NULL;

  Conversation_Store *p_store = calloc(1, sizeof(*p_store));
  if (!p_store)
    return NULL;
  p_store->p_connection = Deita_Connection_Create(
      DEITA_DATABASE_TYPE_SQLITE3,
      database_path);
  if (!p_store->p_connection ||
      !Deita_Connection_Is_Open(p_store->p_connection))
  {
    if (p_store->p_connection)
      Deita_Connection_Close(p_store->p_connection);
    free(p_store);
    return NULL;
  }
  if (pthread_mutex_init(&p_store->mutex, NULL) != 0)
  {
    Deita_Connection_Close(p_store->p_connection);
    free(p_store);
    return NULL;
  }

  const char *schema =
      "PRAGMA foreign_keys = ON;"
      "PRAGMA journal_mode = WAL;"
      "CREATE TABLE IF NOT EXISTS conversations ("
      "id TEXT PRIMARY KEY,"
      "copilot_session_id TEXT NOT NULL UNIQUE,"
      "title TEXT NOT NULL DEFAULT '',"
      "status TEXT NOT NULL DEFAULT 'active',"
      "created_at INTEGER NOT NULL DEFAULT (strftime('%s','now')),"
      "updated_at INTEGER NOT NULL DEFAULT (strftime('%s','now'))"
      ");"
      "CREATE TABLE IF NOT EXISTS conversation_turns ("
      "id INTEGER PRIMARY KEY AUTOINCREMENT,"
      "conversation_id TEXT NOT NULL,"
      "sequence INTEGER NOT NULL,"
      "role TEXT NOT NULL,"
      "content TEXT NOT NULL DEFAULT '',"
      "status TEXT NOT NULL,"
      "request_id TEXT,"
      "error_message TEXT,"
      "input_tokens INTEGER NOT NULL DEFAULT 0,"
      "output_tokens INTEGER NOT NULL DEFAULT 0,"
      "created_at INTEGER NOT NULL DEFAULT (strftime('%s','now')),"
      "completed_at INTEGER,"
      "FOREIGN KEY(conversation_id) REFERENCES conversations(id) "
      "ON DELETE CASCADE,"
      "UNIQUE(conversation_id, sequence)"
      ");"
      "CREATE UNIQUE INDEX IF NOT EXISTS "
      "idx_conversation_active_request "
      "ON conversation_turns(conversation_id) "
      "WHERE role = 'assistant' AND status = 'active';"
      "CREATE INDEX IF NOT EXISTS idx_conversations_updated "
      "ON conversations(updated_at DESC);"
      "CREATE INDEX IF NOT EXISTS idx_turns_conversation_sequence "
      "ON conversation_turns(conversation_id, sequence);";
  if (Deita_Query_Execute_Update(p_store->p_connection, schema) < 0)
  {
    Conversation_Store_Destroy(p_store);
    return NULL;
  }
  if (Deita_Query_Execute_Update(
          p_store->p_connection,
          "UPDATE conversation_turns "
          "SET status = 'failed', "
          "error_message = 'Interrupted by server restart', "
          "completed_at = strftime('%s','now') "
          "WHERE role = 'assistant' AND status = 'active'") < 0)
  {
    Conversation_Store_Destroy(p_store);
    return NULL;
  }
  return p_store;
}

void Conversation_Store_Destroy(Conversation_Store *p_store)
{
  if (!p_store)
    return;
  if (p_store->p_connection)
    Deita_Connection_Close(p_store->p_connection);
  pthread_mutex_destroy(&p_store->mutex);
  free(p_store);
}

Conversation_Store_Result Conversation_Store_Create_Conversation(
    Conversation_Store *p_store,
    const char *title,
    char output_id[37])
{
  if (!p_store || !output_id ||
      !Conversation_Store_Generate_UUID(output_id))
    return CONVERSATION_STORE_ERROR;

  const char *parameters[] = {
    output_id,
    output_id,
    title ? title : "",
  };
  pthread_mutex_lock(&p_store->mutex);
  int32 result = Deita_Query_Execute_Update_Prepared(
      p_store->p_connection,
      "INSERT INTO conversations (id, copilot_session_id, title) "
      "VALUES (?, ?, ?)",
      3,
      parameters);
  pthread_mutex_unlock(&p_store->mutex);
  return result < 0 ? CONVERSATION_STORE_ERROR : CONVERSATION_STORE_OK;
}

Conversation_Store_Result Conversation_Store_Get(
    Conversation_Store *p_store,
    const char *conversation_id,
    Conversation_Record *p_record,
    Dowa_Arena *p_arena)
{
  if (!p_store || !conversation_id || !p_record || !p_arena)
    return CONVERSATION_STORE_ERROR;
  memset(p_record, 0, sizeof(*p_record));

  const char *parameters[] = {conversation_id};
  pthread_mutex_lock(&p_store->mutex);
  Deita_Result_Set *p_result = Deita_Query_Execute_Prepared(
      p_store->p_connection,
      "SELECT id, copilot_session_id, title, status, created_at, updated_at "
      "FROM conversations WHERE id = ? AND status != 'deleted'",
      1,
      parameters,
      p_arena);
  if (!p_result || !Deita_Result_Set_Next(p_result))
  {
    if (p_result)
      Deita_Result_Set_Free(p_result);
    pthread_mutex_unlock(&p_store->mutex);
    return CONVERSATION_STORE_NOT_FOUND;
  }
  p_record->id = Conversation_Store_Copy_Text(
      Deita_Result_Set_Get_Text(p_result, 0), p_arena);
  p_record->copilot_session_id = Conversation_Store_Copy_Text(
      Deita_Result_Set_Get_Text(p_result, 1), p_arena);
  p_record->title = Conversation_Store_Copy_Text(
      Deita_Result_Set_Get_Text(p_result, 2), p_arena);
  p_record->status = Conversation_Store_Copy_Text(
      Deita_Result_Set_Get_Text(p_result, 3), p_arena);
  p_record->created_at = Deita_Result_Set_Get_Integer(p_result, 4);
  p_record->updated_at = Deita_Result_Set_Get_Integer(p_result, 5);
  Deita_Result_Set_Free(p_result);

  p_result = Deita_Query_Execute_Prepared(
      p_store->p_connection,
      "SELECT id, sequence, role, content, status, request_id, "
      "error_message, input_tokens, output_tokens, created_at, completed_at "
      "FROM (SELECT id, sequence, role, content, status, request_id, "
      "error_message, input_tokens, output_tokens, created_at, completed_at "
      "FROM conversation_turns WHERE conversation_id = ? "
      "ORDER BY sequence DESC LIMIT 20) ORDER BY sequence",
      1,
      parameters,
      p_arena);
  if (!p_result)
  {
    pthread_mutex_unlock(&p_store->mutex);
    return CONVERSATION_STORE_ERROR;
  }
  while (p_result && Deita_Result_Set_Next(p_result))
  {
    Conversation_Turn turn = {0};
    turn.id = Deita_Result_Set_Get_Integer(p_result, 0);
    turn.sequence = Deita_Result_Set_Get_Integer(p_result, 1);
    turn.role = Conversation_Store_Copy_Text(
        Deita_Result_Set_Get_Text(p_result, 2), p_arena);
    turn.content = Conversation_Store_Copy_Text(
        Deita_Result_Set_Get_Text(p_result, 3), p_arena);
    turn.status = Conversation_Store_Copy_Text(
        Deita_Result_Set_Get_Text(p_result, 4), p_arena);
    turn.request_id = Conversation_Store_Copy_Text(
        Deita_Result_Set_Get_Text(p_result, 5), p_arena);
    turn.error_message = Conversation_Store_Copy_Text(
        Deita_Result_Set_Get_Text(p_result, 6), p_arena);
    turn.input_tokens = Deita_Result_Set_Get_Integer(p_result, 7);
    turn.output_tokens = Deita_Result_Set_Get_Integer(p_result, 8);
    turn.created_at = Deita_Result_Set_Get_Integer(p_result, 9);
    turn.completed_at = Deita_Result_Set_Get_Integer(p_result, 10);
    Dowa_Array_Push_Arena(p_record->turns, turn, p_arena);
  }
  boolean turns_error = Deita_Result_Set_Has_Error(p_result);
  if (p_result)
    Deita_Result_Set_Free(p_result);
  pthread_mutex_unlock(&p_store->mutex);
  return turns_error ? CONVERSATION_STORE_ERROR : CONVERSATION_STORE_OK;
}

Conversation_Store_Result Conversation_Store_Update_Title(
    Conversation_Store *p_store,
    const char *conversation_id,
    const char *title)
{
  if (!p_store || !conversation_id || !title)
    return CONVERSATION_STORE_ERROR;
  const char *parameters[] = {title, conversation_id};
  pthread_mutex_lock(&p_store->mutex);
  int32 result = Deita_Query_Execute_Update_Prepared(
      p_store->p_connection,
      "UPDATE conversations SET title = ?, "
      "updated_at = strftime('%s','now') "
      "WHERE id = ? AND status != 'deleted'",
      2,
      parameters);
  pthread_mutex_unlock(&p_store->mutex);
  if (result < 0)
    return CONVERSATION_STORE_ERROR;
  return result == 0 ? CONVERSATION_STORE_NOT_FOUND : CONVERSATION_STORE_OK;
}

Conversation_Store_Result Conversation_Store_Delete(
    Conversation_Store *p_store,
    const char *conversation_id)
{
  if (!p_store || !conversation_id)
    return CONVERSATION_STORE_ERROR;
  const char *parameters[] = {conversation_id};
  pthread_mutex_lock(&p_store->mutex);
  int32 result = Deita_Query_Execute_Update_Prepared(
      p_store->p_connection,
      "DELETE FROM conversations WHERE id = ?",
      1,
      parameters);
  pthread_mutex_unlock(&p_store->mutex);
  if (result < 0)
    return CONVERSATION_STORE_ERROR;
  return result == 0 ? CONVERSATION_STORE_NOT_FOUND : CONVERSATION_STORE_OK;
}

Conversation_Store_Result Conversation_Store_Begin_Turn(
    Conversation_Store *p_store,
    const char *conversation_id,
    const char *request_id,
    const char *prompt)
{
  if (!p_store || !conversation_id || !request_id || !prompt)
    return CONVERSATION_STORE_ERROR;

  pthread_mutex_lock(&p_store->mutex);
  if (Deita_Query_Execute_Update(
          p_store->p_connection, "BEGIN IMMEDIATE") < 0)
  {
    pthread_mutex_unlock(&p_store->mutex);
    return CONVERSATION_STORE_ERROR;
  }
  if (!Conversation_Store_Exists_Locked(p_store, conversation_id))
  {
    Conversation_Store_Rollback(
        p_store, CONVERSATION_STORE_NOT_FOUND);
    pthread_mutex_unlock(&p_store->mutex);
    return CONVERSATION_STORE_NOT_FOUND;
  }

  Dowa_Arena *p_arena = Dowa_Arena_Create(2048);
  const char *parameters[] = {conversation_id};
  Deita_Result_Set *p_result = Deita_Query_Execute_Prepared(
      p_store->p_connection,
      "SELECT COALESCE(MAX(sequence), 0), "
      "SUM(CASE WHEN role = 'assistant' AND status = 'active' "
      "THEN 1 ELSE 0 END) "
      "FROM conversation_turns WHERE conversation_id = ?",
      1,
      parameters,
      p_arena);
  if (!p_result || !Deita_Result_Set_Next(p_result))
  {
    if (p_result)
      Deita_Result_Set_Free(p_result);
    Dowa_Arena_Free(p_arena);
    Conversation_Store_Rollback(p_store, CONVERSATION_STORE_ERROR);
    pthread_mutex_unlock(&p_store->mutex);
    return CONVERSATION_STORE_ERROR;
  }
  int64 next_sequence = Deita_Result_Set_Get_Integer(p_result, 0) + 1;
  int64 active_count = Deita_Result_Set_Get_Integer(p_result, 1);
  Deita_Result_Set_Free(p_result);
  Dowa_Arena_Free(p_arena);
  if (active_count > 0)
  {
    Conversation_Store_Rollback(p_store, CONVERSATION_STORE_CONFLICT);
    pthread_mutex_unlock(&p_store->mutex);
    return CONVERSATION_STORE_CONFLICT;
  }

  char user_sequence[32];
  char assistant_sequence[32];
  snprintf(user_sequence, sizeof(user_sequence), "%lld",
           (long long)next_sequence);
  snprintf(assistant_sequence, sizeof(assistant_sequence), "%lld",
           (long long)(next_sequence + 1));
  const char *user_parameters[] = {
    conversation_id, user_sequence, prompt, request_id,
  };
  if (Deita_Query_Execute_Update_Prepared(
          p_store->p_connection,
          "INSERT INTO conversation_turns "
          "(conversation_id, sequence, role, content, status, request_id, "
          "completed_at) VALUES (?, ?, 'user', ?, 'complete', ?, "
          "strftime('%s','now'))",
          4,
          user_parameters) < 0)
  {
    Conversation_Store_Rollback(p_store, CONVERSATION_STORE_ERROR);
    pthread_mutex_unlock(&p_store->mutex);
    return CONVERSATION_STORE_ERROR;
  }
  const char *assistant_parameters[] = {
    conversation_id, assistant_sequence, request_id,
  };
  if (Deita_Query_Execute_Update_Prepared(
          p_store->p_connection,
          "INSERT INTO conversation_turns "
          "(conversation_id, sequence, role, status, request_id) "
          "VALUES (?, ?, 'assistant', 'active', ?)",
          3,
          assistant_parameters) < 0 ||
      Deita_Query_Execute_Update_Prepared(
          p_store->p_connection,
          "UPDATE conversations SET updated_at = strftime('%s','now') "
          "WHERE id = ?",
          1,
          parameters) < 0 ||
      Deita_Query_Execute_Update(
          p_store->p_connection, "COMMIT") < 0)
  {
    Conversation_Store_Rollback(p_store, CONVERSATION_STORE_ERROR);
    pthread_mutex_unlock(&p_store->mutex);
    return CONVERSATION_STORE_ERROR;
  }
  pthread_mutex_unlock(&p_store->mutex);
  return CONVERSATION_STORE_OK;
}

Conversation_Store_Result Conversation_Store_Complete_Turn(
    Conversation_Store *p_store,
    const char *conversation_id,
    const char *request_id,
    const char *content,
    int64 input_tokens,
    int64 output_tokens)
{
  if (!p_store || !conversation_id || !request_id || !content)
    return CONVERSATION_STORE_ERROR;
  char input_value[32];
  char output_value[32];
  snprintf(input_value, sizeof(input_value), "%lld", (long long)input_tokens);
  snprintf(output_value, sizeof(output_value), "%lld", (long long)output_tokens);
  const char *parameters[] = {
    content, input_value, output_value, conversation_id, request_id,
  };
  pthread_mutex_lock(&p_store->mutex);
  int32 result = Deita_Query_Execute_Update_Prepared(
      p_store->p_connection,
      "UPDATE conversation_turns SET content = ?, status = 'complete', "
      "input_tokens = ?, output_tokens = ?, "
      "completed_at = strftime('%s','now') "
      "WHERE conversation_id = ? AND request_id = ? "
      "AND role = 'assistant' AND status = 'active'",
      5,
      parameters);
  pthread_mutex_unlock(&p_store->mutex);
  if (result < 0)
    return CONVERSATION_STORE_ERROR;
  return result == 0 ? CONVERSATION_STORE_NOT_FOUND : CONVERSATION_STORE_OK;
}

Conversation_Store_Result Conversation_Store_Fail_Turn(
    Conversation_Store *p_store,
    const char *conversation_id,
    const char *request_id,
    const char *error_message,
    boolean aborted)
{
  if (!p_store || !conversation_id || !request_id)
    return CONVERSATION_STORE_ERROR;
  const char *parameters[] = {
    aborted ? "aborted" : "failed",
    error_message ? error_message : "",
    conversation_id,
    request_id,
  };
  pthread_mutex_lock(&p_store->mutex);
  int32 result = Deita_Query_Execute_Update_Prepared(
      p_store->p_connection,
      "UPDATE conversation_turns SET status = ?, error_message = ?, "
      "completed_at = strftime('%s','now') "
      "WHERE conversation_id = ? AND request_id = ? "
      "AND role = 'assistant' AND status = 'active'",
      4,
      parameters);
  pthread_mutex_unlock(&p_store->mutex);
  if (result < 0)
    return CONVERSATION_STORE_ERROR;
  return result == 0 ? CONVERSATION_STORE_NOT_FOUND : CONVERSATION_STORE_OK;
}