#include "amduat/pel/queue.h"

#include "amduat/hash/asl1.h"

#include <errno.h>
#include <fcntl.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>

static void 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 void free_ref(amduat_reference_t *ref) {
  if (ref == NULL) {
    return;
  }
  free((void *)ref->digest.data);
  ref->digest.data = NULL;
  ref->digest.len = 0;
}

static bool make_ref(amduat_hash_id_t hash_id,
                     uint8_t fill,
                     amduat_reference_t *out_ref) {
  const amduat_hash_asl1_desc_t *desc;
  uint8_t *digest;

  if (out_ref == NULL) {
    return false;
  }
  desc = amduat_hash_asl1_desc_lookup(hash_id);
  if (desc == NULL || desc->digest_len == 0) {
    return false;
  }
  digest = (uint8_t *)malloc(desc->digest_len);
  if (digest == NULL) {
    return false;
  }
  memset(digest, fill, desc->digest_len);
  *out_ref = amduat_reference(hash_id, amduat_octets(digest, desc->digest_len));
  return true;
}

static bool 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 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 char *make_temp_root(void) {
  char *templ;
  const char template_prefix[] = "amduat_test_pel_queue_XXXXXX";

  templ = (char *)malloc(sizeof(template_prefix));
  if (templ == NULL) {
    return NULL;
  }
  memcpy(templ, template_prefix, sizeof(template_prefix));
  if (mkdtemp(templ) == NULL) {
    free(templ);
    return NULL;
  }
  return templ;
}

static void cleanup_root(const char *root) {
  char *path = NULL;

  if (root == NULL) {
    return;
  }

  if (join_path(root, "PEL_QUEUE", &path)) {
    unlink(path);
    free(path);
  }
  if (join_path(root, "PEL_HEAD", &path)) {
    unlink(path);
    free(path);
  }
  if (join_path(root, "TGK_HEAD", &path)) {
    unlink(path);
    free(path);
  }
  rmdir(root);
}

static int test_append_and_read(void) {
  char *root = NULL;
  amduat_reference_t ref1;
  amduat_reference_t ref2;
  amduat_reference_t out_ref;
  uint64_t offset = 0;
  uint64_t next_offset = 0;
  int exit_code = 1;

  root = make_temp_root();
  if (root == NULL) {
    fprintf(stderr, "temp root failed\n");
    return exit_code;
  }

  if (!make_ref(AMDUAT_HASH_ASL1_ID_SHA256, 0x11u, &ref1) ||
      !make_ref(AMDUAT_HASH_ASL1_ID_SHA256, 0x22u, &ref2)) {
    fprintf(stderr, "ref alloc failed\n");
    goto cleanup;
  }

  if (!amduat_pel_queue_append_trace_ref(root, ref1) ||
      !amduat_pel_queue_append_trace_ref(root, ref2)) {
    fprintf(stderr, "append failed\n");
    goto cleanup;
  }

  if (!amduat_pel_queue_read_next(root, offset, &out_ref, &next_offset)) {
    fprintf(stderr, "read first failed\n");
    goto cleanup;
  }
  if (!amduat_reference_eq(out_ref, ref1)) {
    fprintf(stderr, "first ref mismatch\n");
    free_ref(&out_ref);
    goto cleanup;
  }
  free_ref(&out_ref);
  offset = next_offset;

  if (!amduat_pel_queue_read_next(root, offset, &out_ref, &next_offset)) {
    fprintf(stderr, "read second failed\n");
    goto cleanup;
  }
  if (!amduat_reference_eq(out_ref, ref2)) {
    fprintf(stderr, "second ref mismatch\n");
    free_ref(&out_ref);
    goto cleanup;
  }
  free_ref(&out_ref);
  offset = next_offset;

  if (amduat_pel_queue_read_next(root, offset, &out_ref, &next_offset)) {
    fprintf(stderr, "expected end of queue\n");
    free_ref(&out_ref);
    goto cleanup;
  }
  if (next_offset != offset) {
    fprintf(stderr, "unexpected offset on empty read\n");
    goto cleanup;
  }

  exit_code = 0;

cleanup:
  free_ref(&ref1);
  free_ref(&ref2);
  cleanup_root(root);
  free(root);
  return exit_code;
}

static int test_truncated_tail(void) {
  char *root = NULL;
  char *queue_path = NULL;
  amduat_reference_t ref;
  amduat_reference_t out_ref;
  uint64_t offset = 0;
  uint64_t next_offset = 0;
  uint8_t len_bytes[4];
  int fd;
  int exit_code = 1;

  root = make_temp_root();
  if (root == NULL) {
    fprintf(stderr, "temp root failed\n");
    return exit_code;
  }

  if (!make_ref(AMDUAT_HASH_ASL1_ID_SHA256, 0x33u, &ref)) {
    fprintf(stderr, "ref alloc failed\n");
    goto cleanup;
  }

  if (!amduat_pel_queue_append_trace_ref(root, ref)) {
    fprintf(stderr, "append failed\n");
    goto cleanup;
  }

  if (!join_path(root, "PEL_QUEUE", &queue_path)) {
    fprintf(stderr, "queue path failed\n");
    goto cleanup;
  }
  fd = open(queue_path, O_WRONLY | O_APPEND);
  free(queue_path);
  queue_path = NULL;
  if (fd < 0) {
    fprintf(stderr, "open queue failed\n");
    goto cleanup;
  }

  store_u32_be(len_bytes, 7u);
  if (!write_all(fd, len_bytes, sizeof(len_bytes))) {
    close(fd);
    fprintf(stderr, "write truncated failed\n");
    goto cleanup;
  }
  if (close(fd) != 0) {
    fprintf(stderr, "close queue failed\n");
    goto cleanup;
  }

  if (!amduat_pel_queue_read_next(root, offset, &out_ref, &next_offset)) {
    fprintf(stderr, "read first failed\n");
    goto cleanup;
  }
  free_ref(&out_ref);
  offset = next_offset;

  if (amduat_pel_queue_read_next(root, offset, &out_ref, &next_offset)) {
    fprintf(stderr, "expected truncated tail to stop\n");
    free_ref(&out_ref);
    goto cleanup;
  }
  if (next_offset != offset) {
    fprintf(stderr, "unexpected offset on truncated read\n");
    goto cleanup;
  }

  exit_code = 0;

cleanup:
  free_ref(&ref);
  cleanup_root(root);
  free(root);
  return exit_code;
}

static int test_heads_round_trip(void) {
  char *root = NULL;
  uint64_t pel_offset = 12345u;
  uint64_t tgk_offset = 67890u;
  uint64_t loaded = 0u;
  int exit_code = 1;

  root = make_temp_root();
  if (root == NULL) {
    fprintf(stderr, "temp root failed\n");
    return exit_code;
  }

  if (!amduat_pel_queue_store_pel_head(root, pel_offset) ||
      !amduat_pel_queue_load_pel_head(root, &loaded)) {
    fprintf(stderr, "pel head round-trip failed\n");
    goto cleanup;
  }
  if (loaded != pel_offset) {
    fprintf(stderr, "pel head mismatch\n");
    goto cleanup;
  }

  if (!amduat_pel_queue_store_tgk_head(root, tgk_offset) ||
      !amduat_pel_queue_load_tgk_head(root, &loaded)) {
    fprintf(stderr, "tgk head round-trip failed\n");
    goto cleanup;
  }
  if (loaded != tgk_offset) {
    fprintf(stderr, "tgk head mismatch\n");
    goto cleanup;
  }

  exit_code = 0;

cleanup:
  cleanup_root(root);
  free(root);
  return exit_code;
}

int main(void) {
  if (test_append_and_read() != 0) {
    return 1;
  }
  if (test_truncated_tail() != 0) {
    return 1;
  }
  if (test_heads_round_trip() != 0) {
    return 1;
  }
  return 0;
}