#include "amduat/pel/queue.h"

#include "amduat/enc/asl1_core_codec.h"
#include "amduat/hash/asl1.h"

#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

typedef struct {
  const uint8_t *data;
  size_t len;
  size_t offset;
} amduat_cursor_t;

typedef enum {
  AMDUAT_PEL_QUEUE_READ_OK = 0,
  AMDUAT_PEL_QUEUE_READ_NOT_FOUND = 1,
  AMDUAT_PEL_QUEUE_READ_ERR = 2
} amduat_pel_queue_read_status_t;

enum {
  AMDUAT_PEL_QUEUE_RECORD_VERSION = 1,
  AMDUAT_PEL_QUEUE_MIN_RECORD_LEN = 1 + 4 + 2
};

static const char k_pel_queue_path[] = "PEL_QUEUE";
static const char k_pel_head_path[] = "PEL_HEAD";
static const char k_tgk_head_path[] = "TGK_HEAD";
static const char k_pel_head_magic[] = "amduat-pel-head-v1";
static const char k_tgk_head_magic[] = "amduat-tgk-head-v1";

static void amduat_store_u16_be(uint8_t *out, uint16_t value) {
  out[0] = (uint8_t)((value >> 8) & 0xffu);
  out[1] = (uint8_t)(value & 0xffu);
}

static void amduat_store_u32_be(uint8_t *out, uint32_t value) {
  out[0] = (uint8_t)((value >> 24) & 0xffu);
  out[1] = (uint8_t)((value >> 16) & 0xffu);
  out[2] = (uint8_t)((value >> 8) & 0xffu);
  out[3] = (uint8_t)(value & 0xffu);
}

static uint32_t amduat_load_u32_be(const uint8_t *data) {
  return ((uint32_t)data[0] << 24) | ((uint32_t)data[1] << 16) |
         ((uint32_t)data[2] << 8) | (uint32_t)data[3];
}

static bool amduat_read_u8(amduat_cursor_t *cur, uint8_t *out) {
  if (cur->len - cur->offset < 1) {
    return false;
  }
  *out = cur->data[cur->offset];
  cur->offset += 1;
  return true;
}

static bool amduat_read_u32(amduat_cursor_t *cur, uint32_t *out) {
  const uint8_t *data;
  if (cur->len - cur->offset < 4) {
    return false;
  }
  data = cur->data + cur->offset;
  *out = ((uint32_t)data[0] << 24) | ((uint32_t)data[1] << 16) |
         ((uint32_t)data[2] << 8) | (uint32_t)data[3];
  cur->offset += 4;
  return true;
}

static bool amduat_add_size(size_t *acc, size_t add) {
  if (*acc > SIZE_MAX - add) {
    return false;
  }
  *acc += add;
  return true;
}

static bool amduat_reference_bytes_len(amduat_reference_t ref,
                                       size_t *out_len) {
  const amduat_hash_asl1_desc_t *desc;

  if (ref.digest.len != 0 && ref.digest.data == NULL) {
    return false;
  }

  desc = amduat_hash_asl1_desc_lookup(ref.hash_id);
  if (desc == NULL || desc->digest_len == 0) {
    return false;
  }
  if (ref.digest.len != desc->digest_len) {
    return false;
  }

  *out_len = 2 + desc->digest_len;
  return true;
}

static bool amduat_encoded_ref_len(amduat_reference_t ref, size_t *out_len) {
  size_t ref_len;

  if (!amduat_reference_bytes_len(ref, &ref_len)) {
    return false;
  }
  if (ref_len > UINT32_MAX) {
    return false;
  }
  *out_len = 4 + ref_len;
  return true;
}

static bool amduat_write_encoded_ref(uint8_t *buffer,
                                     size_t buffer_len,
                                     size_t *offset,
                                     amduat_reference_t ref) {
  size_t ref_len;

  if (!amduat_reference_bytes_len(ref, &ref_len)) {
    return false;
  }
  if (ref_len > UINT32_MAX) {
    return false;
  }
  if (buffer_len - *offset < 4 + ref_len) {
    return false;
  }

  amduat_store_u32_be(buffer + *offset, (uint32_t)ref_len);
  *offset += 4;
  amduat_store_u16_be(buffer + *offset, ref.hash_id);
  *offset += 2;
  if (ref.digest.len != 0) {
    memcpy(buffer + *offset, ref.digest.data, ref.digest.len);
    *offset += ref.digest.len;
  }

  return true;
}

static bool amduat_read_encoded_ref(amduat_cursor_t *cur,
                                    amduat_reference_t *out_ref) {
  uint32_t ref_len_u32;
  amduat_octets_t ref_bytes;

  if (!amduat_read_u32(cur, &ref_len_u32)) {
    return false;
  }
  if (ref_len_u32 < 2) {
    return false;
  }
  if (cur->len - cur->offset < ref_len_u32) {
    return false;
  }
  ref_bytes = amduat_octets(cur->data + cur->offset, ref_len_u32);
  if (!amduat_enc_asl1_core_decode_reference_v1(ref_bytes, out_ref)) {
    return false;
  }
  cur->offset += ref_len_u32;
  return true;
}

static void amduat_reference_free(amduat_reference_t *ref) {
  if (ref == NULL) {
    return;
  }
  free((void *)ref->digest.data);
  ref->digest.data = NULL;
  ref->digest.len = 0;
}

static bool amduat_join_path(const char *base,
                             const char *segment,
                             char **out_path) {
  size_t base_len;
  size_t seg_len;
  bool needs_sep;
  size_t total_len;
  char *buffer;
  size_t offset;

  if (base == NULL || segment == NULL || out_path == NULL) {
    return false;
  }
  if (base[0] == '\0' || segment[0] == '\0') {
    return false;
  }

  base_len = strlen(base);
  seg_len = strlen(segment);
  needs_sep = base[base_len - 1u] != '/';
  total_len = base_len + (needs_sep ? 1u : 0u) + seg_len + 1u;

  buffer = (char *)malloc(total_len);
  if (buffer == NULL) {
    return false;
  }

  offset = 0u;
  memcpy(buffer + offset, base, base_len);
  offset += base_len;
  if (needs_sep) {
    buffer[offset++] = '/';
  }
  memcpy(buffer + offset, segment, seg_len);
  offset += seg_len;
  buffer[offset] = '\0';

  *out_path = buffer;
  return true;
}

static bool amduat_write_all(int fd, const uint8_t *data, size_t len) {
  size_t total = 0u;

  while (total < len) {
    ssize_t rc = write(fd, data + total, len - total);
    if (rc < 0) {
      if (errno == EINTR) {
        continue;
      }
      return false;
    }
    total += (size_t)rc;
  }

  return true;
}

static bool amduat_read_full(int fd,
                             uint8_t *data,
                             size_t len,
                             size_t *out_read) {
  size_t total = 0u;

  if (out_read == NULL) {
    return false;
  }

  while (total < len) {
    ssize_t rc = read(fd, data + total, len - total);
    if (rc < 0) {
      if (errno == EINTR) {
        continue;
      }
      return false;
    }
    if (rc == 0) {
      break;
    }
    total += (size_t)rc;
  }

  *out_read = total;
  return true;
}

static amduat_pel_queue_read_status_t amduat_read_file(const char *path,
                                                       char **out_bytes,
                                                       size_t *out_size) {
  struct stat st;
  size_t file_size;
  char *buffer;
  size_t total_read;
  int fd;

  if (path == NULL || out_bytes == NULL || out_size == NULL) {
    return AMDUAT_PEL_QUEUE_READ_ERR;
  }
  *out_bytes = NULL;
  *out_size = 0;

  if (stat(path, &st) != 0) {
    if (errno == ENOENT || errno == ENOTDIR) {
      return AMDUAT_PEL_QUEUE_READ_NOT_FOUND;
    }
    return AMDUAT_PEL_QUEUE_READ_ERR;
  }
  if (!S_ISREG(st.st_mode)) {
    return AMDUAT_PEL_QUEUE_READ_ERR;
  }
  if (st.st_size <= 0) {
    return AMDUAT_PEL_QUEUE_READ_ERR;
  }
  if ((uintmax_t)st.st_size > SIZE_MAX) {
    return AMDUAT_PEL_QUEUE_READ_ERR;
  }

  file_size = (size_t)st.st_size;
  buffer = (char *)malloc(file_size + 1u);
  if (buffer == NULL) {
    return AMDUAT_PEL_QUEUE_READ_ERR;
  }

  fd = open(path, O_RDONLY);
  if (fd < 0) {
    free(buffer);
    if (errno == ENOENT || errno == ENOTDIR) {
      return AMDUAT_PEL_QUEUE_READ_NOT_FOUND;
    }
    return AMDUAT_PEL_QUEUE_READ_ERR;
  }

  total_read = 0u;
  while (total_read < file_size) {
    ssize_t rc = read(fd, buffer + total_read, file_size - total_read);
    if (rc < 0) {
      if (errno == EINTR) {
        continue;
      }
      close(fd);
      free(buffer);
      return AMDUAT_PEL_QUEUE_READ_ERR;
    }
    if (rc == 0) {
      close(fd);
      free(buffer);
      return AMDUAT_PEL_QUEUE_READ_ERR;
    }
    total_read += (size_t)rc;
  }

  if (close(fd) != 0) {
    free(buffer);
    return AMDUAT_PEL_QUEUE_READ_ERR;
  }

  buffer[file_size] = '\0';
  *out_bytes = buffer;
  *out_size = file_size;
  return AMDUAT_PEL_QUEUE_READ_OK;
}

static bool amduat_parse_head_text(const char *text,
                                   const char *magic,
                                   uint64_t *out_offset) {
  size_t magic_len;
  const char *cursor;
  char *end;
  unsigned long long value;

  if (text == NULL || magic == NULL || out_offset == NULL) {
    return false;
  }

  magic_len = strlen(magic);
  if (strncmp(text, magic, magic_len) != 0) {
    return false;
  }
  if (text[magic_len] != '\n') {
    return false;
  }

  cursor = text + magic_len + 1;
  if (strncmp(cursor, "offset=", 7) != 0) {
    return false;
  }
  cursor += 7;
  if (*cursor == '\0') {
    return false;
  }
  errno = 0;
  value = strtoull(cursor, &end, 10);
  if (errno != 0 || end == cursor) {
    return false;
  }
  if (*end == '\n') {
    end++;
  }
  if (*end != '\0') {
    return false;
  }
  if (value > (unsigned long long)UINT64_MAX) {
    return false;
  }

  *out_offset = (uint64_t)value;
  return true;
}

static bool amduat_load_head(const char *root,
                             const char *path_segment,
                             const char *magic,
                             uint64_t *out_offset) {
  char *path = NULL;
  amduat_pel_queue_read_status_t read_status;
  char *head_text = NULL;
  size_t head_size = 0u;
  bool ok;

  if (root == NULL || root[0] == '\0' || out_offset == NULL) {
    return false;
  }

  if (!amduat_join_path(root, path_segment, &path)) {
    return false;
  }

  read_status = amduat_read_file(path, &head_text, &head_size);
  free(path);
  (void)head_size;
  if (read_status == AMDUAT_PEL_QUEUE_READ_NOT_FOUND) {
    *out_offset = 0;
    return true;
  }
  if (read_status != AMDUAT_PEL_QUEUE_READ_OK) {
    return false;
  }

  ok = amduat_parse_head_text(head_text, magic, out_offset);
  free(head_text);
  return ok;
}

static bool amduat_store_head(const char *root,
                              const char *path_segment,
                              const char *magic,
                              uint64_t offset) {
  char *path = NULL;
  char buffer[128];
  int buffer_len;
  int fd;
  bool ok;

  if (root == NULL || root[0] == '\0') {
    return false;
  }

  if (!amduat_join_path(root, path_segment, &path)) {
    return false;
  }

  buffer_len = snprintf(buffer,
                        sizeof(buffer),
                        "%s\noffset=%" PRIu64 "\n",
                        magic,
                        offset);
  if (buffer_len <= 0 || (size_t)buffer_len >= sizeof(buffer)) {
    free(path);
    return false;
  }

  fd = open(path, O_WRONLY | O_CREAT | O_TRUNC, 0644);
  free(path);
  if (fd < 0) {
    return false;
  }

  ok = amduat_write_all(fd, (const uint8_t *)buffer, (size_t)buffer_len);
  if (!ok) {
    close(fd);
    return false;
  }
  if (close(fd) != 0) {
    return false;
  }

  return true;
}

bool amduat_pel_queue_append_trace_ref(const char *root,
                                       amduat_reference_t ref) {
  char *path = NULL;
  size_t encoded_len;
  size_t record_len;
  size_t total_len;
  uint8_t *buffer;
  size_t offset;
  int fd;
  bool ok;

  if (root == NULL || root[0] == '\0') {
    return false;
  }

  if (!amduat_encoded_ref_len(ref, &encoded_len)) {
    return false;
  }
  record_len = 1u;
  if (!amduat_add_size(&record_len, encoded_len)) {
    return false;
  }
  if (record_len > UINT32_MAX) {
    return false;
  }
  total_len = 4u;
  if (!amduat_add_size(&total_len, record_len)) {
    return false;
  }

  buffer = (uint8_t *)malloc(total_len);
  if (buffer == NULL) {
    return false;
  }

  amduat_store_u32_be(buffer, (uint32_t)record_len);
  offset = 4u;
  buffer[offset++] = AMDUAT_PEL_QUEUE_RECORD_VERSION;
  if (!amduat_write_encoded_ref(buffer, total_len, &offset, ref)) {
    free(buffer);
    return false;
  }
  if (offset != total_len) {
    free(buffer);
    return false;
  }

  if (!amduat_join_path(root, k_pel_queue_path, &path)) {
    free(buffer);
    return false;
  }

  fd = open(path, O_WRONLY | O_APPEND | O_CREAT, 0644);
  free(path);
  if (fd < 0) {
    free(buffer);
    return false;
  }

  ok = amduat_write_all(fd, buffer, total_len);
  free(buffer);
  if (!ok) {
    close(fd);
    return false;
  }
  if (close(fd) != 0) {
    return false;
  }

  return true;
}

bool amduat_pel_queue_read_next(const char *root,
                                uint64_t offset,
                                amduat_reference_t *out_ref,
                                uint64_t *out_next_offset) {
  char *path = NULL;
  int fd;
  uint8_t len_bytes[4];
  size_t read_len;
  uint32_t record_len_u32;
  size_t record_len;
  uint8_t *record_bytes;
  amduat_cursor_t cur;
  uint8_t version;
  amduat_reference_t ref;
  uint64_t next_offset;

  if (out_next_offset != NULL) {
    *out_next_offset = offset;
  }

  if (root == NULL || root[0] == '\0' || out_ref == NULL ||
      out_next_offset == NULL) {
    return false;
  }

  if (!amduat_join_path(root, k_pel_queue_path, &path)) {
    return false;
  }

  fd = open(path, O_RDONLY);
  free(path);
  if (fd < 0) {
    if (errno == ENOENT || errno == ENOTDIR) {
      return false;
    }
    return false;
  }

  if (lseek(fd, (off_t)offset, SEEK_SET) == (off_t)-1) {
    close(fd);
    return false;
  }

  if (!amduat_read_full(fd, len_bytes, sizeof(len_bytes), &read_len)) {
    close(fd);
    return false;
  }
  if (read_len < sizeof(len_bytes)) {
    close(fd);
    return false;
  }

  record_len_u32 = amduat_load_u32_be(len_bytes);
  if (record_len_u32 < AMDUAT_PEL_QUEUE_MIN_RECORD_LEN) {
    close(fd);
    return false;
  }
  if (record_len_u32 > SIZE_MAX) {
    close(fd);
    return false;
  }
  record_len = (size_t)record_len_u32;
  record_bytes = (uint8_t *)malloc(record_len);
  if (record_bytes == NULL) {
    close(fd);
    return false;
  }

  if (!amduat_read_full(fd, record_bytes, record_len, &read_len)) {
    free(record_bytes);
    close(fd);
    return false;
  }
  if (read_len < record_len) {
    free(record_bytes);
    close(fd);
    return false;
  }

  if (close(fd) != 0) {
    free(record_bytes);
    return false;
  }

  cur.data = record_bytes;
  cur.len = record_len;
  cur.offset = 0u;
  if (!amduat_read_u8(&cur, &version)) {
    free(record_bytes);
    return false;
  }
  if (version != AMDUAT_PEL_QUEUE_RECORD_VERSION) {
    free(record_bytes);
    return false;
  }

  ref.hash_id = 0;
  ref.digest = amduat_octets(NULL, 0);
  if (!amduat_read_encoded_ref(&cur, &ref)) {
    free(record_bytes);
    return false;
  }
  if (cur.offset != cur.len) {
    amduat_reference_free(&ref);
    free(record_bytes);
    return false;
  }

  if (offset > UINT64_MAX - 4u - record_len_u32) {
    amduat_reference_free(&ref);
    free(record_bytes);
    return false;
  }
  next_offset = offset + 4u + record_len_u32;

  *out_ref = ref;
  *out_next_offset = next_offset;
  free(record_bytes);
  return true;
}

bool amduat_pel_queue_load_pel_head(const char *root, uint64_t *out_offset) {
  return amduat_load_head(root, k_pel_head_path, k_pel_head_magic, out_offset);
}

bool amduat_pel_queue_store_pel_head(const char *root, uint64_t offset) {
  return amduat_store_head(root, k_pel_head_path, k_pel_head_magic, offset);
}

bool amduat_pel_queue_load_tgk_head(const char *root, uint64_t *out_offset) {
  return amduat_load_head(root, k_tgk_head_path, k_tgk_head_magic, out_offset);
}

bool amduat_pel_queue_store_tgk_head(const char *root, uint64_t offset) {
  return amduat_store_head(root, k_tgk_head_path, k_tgk_head_magic, offset);
}