#include "amduat/asl/collection.h"

#include "amduat/enc/asl1_core_codec.h"
#include "amduat/util/log.h"

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

enum {
  AMDUAT_ASL_COLLECTION_MAGIC_LEN = 8,
  AMDUAT_ASL_COLLECTION_VERSION = 1,
  AMDUAT_ASL_COLLECTION_READ_BATCH = 1024u
};

static const uint8_t k_amduat_asl_collection_magic[
    AMDUAT_ASL_COLLECTION_MAGIC_LEN] = {
        'A', 'S', 'L', 'C', 'O', 'L', '1', '\0'
};

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

static void amduat_asl_collection_store_u64_le(uint8_t *out, uint64_t value) {
  out[0] = (uint8_t)(value & 0xffu);
  out[1] = (uint8_t)((value >> 8) & 0xffu);
  out[2] = (uint8_t)((value >> 16) & 0xffu);
  out[3] = (uint8_t)((value >> 24) & 0xffu);
  out[4] = (uint8_t)((value >> 32) & 0xffu);
  out[5] = (uint8_t)((value >> 40) & 0xffu);
  out[6] = (uint8_t)((value >> 48) & 0xffu);
  out[7] = (uint8_t)((value >> 56) & 0xffu);
}

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

static bool amduat_asl_collection_build_log_name(const char *name,
                                                 char **out_name) {
  size_t name_len;
  size_t total_len;
  char *buffer;
  size_t offset;

  if (name == NULL || out_name == NULL) {
    return false;
  }
  if (!amduat_asl_pointer_name_is_valid(name)) {
    return false;
  }
  name_len = strlen(name);
  total_len = 11u + name_len + 4u + 1u;
  buffer = (char *)malloc(total_len);
  if (buffer == NULL) {
    return false;
  }
  offset = 0u;
  memcpy(buffer + offset, "collection/", 11u);
  offset += 11u;
  memcpy(buffer + offset, name, name_len);
  offset += name_len;
  memcpy(buffer + offset, "/log", 4u);
  offset += 4u;
  buffer[offset] = '\0';
  *out_name = buffer;
  return true;
}

static bool amduat_asl_collection_build_head_name(const char *name,
                                                  char **out_name) {
  size_t name_len;
  size_t total_len;
  char *buffer;
  size_t offset;

  if (name == NULL || out_name == NULL) {
    return false;
  }
  if (!amduat_asl_pointer_name_is_valid(name)) {
    return false;
  }
  name_len = strlen(name);
  total_len = 11u + name_len + 5u + 1u;
  buffer = (char *)malloc(total_len);
  if (buffer == NULL) {
    return false;
  }
  offset = 0u;
  memcpy(buffer + offset, "collection/", 11u);
  offset += 11u;
  memcpy(buffer + offset, name, name_len);
  offset += name_len;
  memcpy(buffer + offset, "/head", 5u);
  offset += 5u;
  buffer[offset] = '\0';
  *out_name = buffer;
  return true;
}

static bool amduat_asl_collection_encode_snapshot_payload(
    uint64_t snapshot_offset,
    const amduat_reference_t *refs,
    size_t refs_len,
    amduat_octets_t *out_payload) {
  size_t total_len = 0u;
  uint8_t *buffer;
  size_t offset = 0u;

  if (out_payload == NULL) {
    return false;
  }
  out_payload->data = NULL;
  out_payload->len = 0u;

  if (refs_len > UINT32_MAX) {
    return false;
  }

  if (!amduat_asl_collection_add_size(
          &total_len, AMDUAT_ASL_COLLECTION_MAGIC_LEN + 4u + 8u + 4u)) {
    return false;
  }
  for (size_t i = 0u; i < refs_len; ++i) {
    amduat_octets_t ref_bytes = amduat_octets(NULL, 0u);
    if (!amduat_enc_asl1_core_encode_reference_v1(refs[i], &ref_bytes)) {
      return false;
    }
    if (!amduat_asl_collection_add_size(&total_len, 4u + ref_bytes.len)) {
      free((void *)ref_bytes.data);
      return false;
    }
    free((void *)ref_bytes.data);
  }

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

  memcpy(buffer + offset, k_amduat_asl_collection_magic,
         AMDUAT_ASL_COLLECTION_MAGIC_LEN);
  offset += AMDUAT_ASL_COLLECTION_MAGIC_LEN;
  amduat_asl_collection_store_u32_le(buffer + offset,
                                     AMDUAT_ASL_COLLECTION_VERSION);
  offset += 4u;
  amduat_asl_collection_store_u64_le(buffer + offset, snapshot_offset);
  offset += 8u;
  amduat_asl_collection_store_u32_le(buffer + offset, (uint32_t)refs_len);
  offset += 4u;

  for (size_t i = 0u; i < refs_len; ++i) {
    amduat_octets_t ref_bytes = amduat_octets(NULL, 0u);
    if (!amduat_enc_asl1_core_encode_reference_v1(refs[i], &ref_bytes)) {
      free(buffer);
      return false;
    }
    amduat_asl_collection_store_u32_le(buffer + offset,
                                       (uint32_t)ref_bytes.len);
    offset += 4u;
    memcpy(buffer + offset, ref_bytes.data, ref_bytes.len);
    offset += ref_bytes.len;
    free((void *)ref_bytes.data);
  }

  if (offset != total_len) {
    free(buffer);
    return false;
  }
  *out_payload = amduat_octets(buffer, total_len);
  return true;
}

bool amduat_asl_collection_store_init(
    amduat_asl_collection_store_t *collection_store,
    const char *root_path,
    amduat_asl_store_t *store) {
  if (collection_store == NULL || root_path == NULL || store == NULL) {
    return false;
  }
  collection_store->store = store;
  if (!amduat_asl_pointer_store_init(&collection_store->pointer_store,
                                     root_path)) {
    return false;
  }
  if (!amduat_asl_log_store_init(&collection_store->log_store, root_path,
                                 store, &collection_store->pointer_store)) {
    return false;
  }
  return true;
}

amduat_asl_collection_error_t amduat_asl_collection_append(
    amduat_asl_collection_store_t *collection_store,
    const char *collection_name,
    amduat_reference_t record_ref,
    uint16_t kind,
    amduat_octets_t actor,
    uint64_t *out_offset) {
  char *log_name = NULL;
  amduat_asl_log_entry_t entry;
  amduat_asl_store_error_t err;

  if (collection_store == NULL || collection_store->store == NULL ||
      collection_name == NULL) {
    return AMDUAT_ASL_COLLECTION_ERR_INTEGRITY;
  }
  if (!amduat_asl_pointer_name_is_valid(collection_name)) {
    return AMDUAT_ASL_COLLECTION_ERR_INVALID_NAME;
  }
  if (!amduat_asl_collection_build_log_name(collection_name, &log_name)) {
    return AMDUAT_ASL_COLLECTION_ERR_IO;
  }

  memset(&entry, 0, sizeof(entry));
  entry.kind = kind;
  entry.payload_ref = record_ref;
  entry.has_actor = actor.len != 0u;
  entry.actor = actor;

  err = amduat_asl_log_append(&collection_store->log_store, log_name,
                              &entry, 1u, out_offset);
  free(log_name);
  if (err != AMDUAT_ASL_STORE_OK) {
    return AMDUAT_ASL_COLLECTION_ERR_IO;
  }
  return AMDUAT_ASL_COLLECTION_OK;
}

amduat_asl_collection_error_t amduat_asl_collection_snapshot(
    amduat_asl_collection_store_t *collection_store,
    const char *collection_name,
    uint64_t up_to_offset,
    amduat_reference_t *out_snapshot_ref,
    bool *out_swapped) {
  char *log_name = NULL;
  char *head_name = NULL;
  amduat_reference_t *refs = NULL;
  size_t refs_len = 0u;
  uint64_t from = 0u;
  bool end = false;
  uint64_t next_offset = 0u;
  amduat_asl_record_t snapshot_record;
  amduat_octets_t payload = amduat_octets(NULL, 0u);
  amduat_reference_t snapshot_ref;
  amduat_asl_pointer_error_t ptr_err;
  bool head_exists = false;
  amduat_reference_t head_ref;
  bool swapped = false;

  if (out_snapshot_ref != NULL) {
    out_snapshot_ref->hash_id = 0u;
    out_snapshot_ref->digest = amduat_octets(NULL, 0u);
  }
  if (out_swapped != NULL) {
    *out_swapped = false;
  }
  if (collection_store == NULL || collection_store->store == NULL ||
      collection_name == NULL) {
    return AMDUAT_ASL_COLLECTION_ERR_INTEGRITY;
  }
  if (!amduat_asl_pointer_name_is_valid(collection_name)) {
    return AMDUAT_ASL_COLLECTION_ERR_INVALID_NAME;
  }

  if (!amduat_asl_collection_build_log_name(collection_name, &log_name) ||
      !amduat_asl_collection_build_head_name(collection_name, &head_name)) {
    free(log_name);
    free(head_name);
    return AMDUAT_ASL_COLLECTION_ERR_IO;
  }

  while (!end && (up_to_offset == UINT64_MAX || from <= up_to_offset)) {
    size_t limit = AMDUAT_ASL_COLLECTION_READ_BATCH;
    amduat_asl_log_entry_t *entries = NULL;
    size_t entries_len = 0u;
    amduat_asl_store_error_t err = amduat_asl_log_read(
        &collection_store->log_store, log_name, from, limit, &entries,
        &entries_len, &next_offset, &end);
    if (err != AMDUAT_ASL_STORE_OK) {
      free(log_name);
      free(head_name);
      amduat_asl_collection_refs_free(refs, refs_len);
      return AMDUAT_ASL_COLLECTION_ERR_IO;
    }
    for (size_t i = 0u; i < entries_len; ++i) {
      uint64_t offset = from + i;
      if (up_to_offset != UINT64_MAX && offset > up_to_offset) {
        end = true;
        break;
      }
      amduat_reference_t *next =
          (amduat_reference_t *)realloc(
              refs, (refs_len + 1u) * sizeof(*refs));
      if (next == NULL) {
        amduat_asl_log_entries_free(entries, entries_len);
        amduat_asl_collection_refs_free(refs, refs_len);
        free(log_name);
        free(head_name);
        return AMDUAT_ASL_COLLECTION_ERR_IO;
      }
      refs = next;
      refs[refs_len] = entries[i].payload_ref;
      if (!amduat_octets_clone(entries[i].payload_ref.digest,
                               &refs[refs_len].digest)) {
        amduat_asl_log_entries_free(entries, entries_len);
        amduat_asl_collection_refs_free(refs, refs_len);
        free(log_name);
        free(head_name);
        return AMDUAT_ASL_COLLECTION_ERR_IO;
      }
      refs_len++;
    }
    amduat_asl_log_entries_free(entries, entries_len);
    from = next_offset;
    if (entries_len == 0u) {
      break;
    }
  }

  {
    uint64_t snapshot_offset = from;
    if (up_to_offset != UINT64_MAX &&
        up_to_offset < UINT64_MAX &&
        snapshot_offset > up_to_offset + 1u) {
      snapshot_offset = up_to_offset + 1u;
    }
    if (!amduat_asl_collection_encode_snapshot_payload(
            snapshot_offset, refs, refs_len, &payload)) {
      amduat_asl_collection_refs_free(refs, refs_len);
      free(log_name);
      free(head_name);
      return AMDUAT_ASL_COLLECTION_ERR_INTEGRITY;
    }
  }

  if (payload.len == 0u || payload.data == NULL) {
    amduat_asl_collection_refs_free(refs, refs_len);
    free(log_name);
    free(head_name);
    return AMDUAT_ASL_COLLECTION_ERR_INTEGRITY;
  }

  {
    amduat_asl_store_error_t err = amduat_asl_record_store_put(
        collection_store->store,
        amduat_octets("collection/snapshot",
                      strlen("collection/snapshot")),
        payload,
        &snapshot_ref);
    amduat_octets_free(&payload);
    if (err != AMDUAT_ASL_STORE_OK) {
      amduat_asl_collection_refs_free(refs, refs_len);
      free(log_name);
      free(head_name);
      return AMDUAT_ASL_COLLECTION_ERR_IO;
    }
  }

  ptr_err = amduat_asl_pointer_get(&collection_store->pointer_store, head_name,
                                   &head_exists, &head_ref);
  if (ptr_err != AMDUAT_ASL_POINTER_OK) {
    amduat_asl_collection_refs_free(refs, refs_len);
    free(log_name);
    free(head_name);
    return AMDUAT_ASL_COLLECTION_ERR_IO;
  }
  ptr_err = amduat_asl_pointer_cas(&collection_store->pointer_store, head_name,
                                   head_exists,
                                   head_exists ? &head_ref : NULL,
                                   &snapshot_ref, &swapped);
  if (head_exists) {
    amduat_reference_free(&head_ref);
  }
  if (ptr_err != AMDUAT_ASL_POINTER_OK) {
    amduat_asl_collection_refs_free(refs, refs_len);
    free(log_name);
    free(head_name);
    return AMDUAT_ASL_COLLECTION_ERR_IO;
  }

  if (out_swapped != NULL) {
    *out_swapped = swapped;
  }
  if (out_snapshot_ref != NULL) {
    *out_snapshot_ref = snapshot_ref;
  }
  if (!swapped) {
    amduat_log(AMDUAT_LOG_DEBUG, "collection snapshot CAS mismatch");
    amduat_asl_collection_refs_free(refs, refs_len);
    free(log_name);
    free(head_name);
    return AMDUAT_ASL_COLLECTION_ERR_CAS_MISMATCH;
  }

  amduat_asl_collection_refs_free(refs, refs_len);
  free(log_name);
  free(head_name);
  return AMDUAT_ASL_COLLECTION_OK;
}

amduat_asl_collection_error_t amduat_asl_collection_read(
    amduat_asl_collection_store_t *collection_store,
    const char *collection_name,
    uint64_t from_offset,
    size_t limit,
    amduat_reference_t **out_record_refs,
    size_t *out_len,
    uint64_t *out_next_offset,
    bool *out_end) {
  char *log_name = NULL;
  amduat_asl_log_entry_t *entries = NULL;
  size_t entries_len = 0u;
  amduat_asl_store_error_t err;

  if (out_record_refs == NULL || out_len == NULL ||
      out_next_offset == NULL || out_end == NULL) {
    return AMDUAT_ASL_COLLECTION_ERR_INTEGRITY;
  }
  *out_record_refs = NULL;
  *out_len = 0u;

  if (collection_store == NULL || collection_store->store == NULL ||
      collection_name == NULL) {
    return AMDUAT_ASL_COLLECTION_ERR_INTEGRITY;
  }
  if (!amduat_asl_pointer_name_is_valid(collection_name)) {
    return AMDUAT_ASL_COLLECTION_ERR_INVALID_NAME;
  }
  if (!amduat_asl_collection_build_log_name(collection_name, &log_name)) {
    return AMDUAT_ASL_COLLECTION_ERR_IO;
  }

  err = amduat_asl_log_read(&collection_store->log_store, log_name,
                            from_offset, limit, &entries, &entries_len,
                            out_next_offset, out_end);
  free(log_name);
  if (err != AMDUAT_ASL_STORE_OK) {
    return AMDUAT_ASL_COLLECTION_ERR_IO;
  }

  if (entries_len != 0u) {
    amduat_reference_t *refs =
        (amduat_reference_t *)calloc(entries_len, sizeof(*refs));
    if (refs == NULL) {
      amduat_asl_log_entries_free(entries, entries_len);
      return AMDUAT_ASL_COLLECTION_ERR_IO;
    }
    for (size_t i = 0u; i < entries_len; ++i) {
      refs[i].hash_id = entries[i].payload_ref.hash_id;
      if (!amduat_octets_clone(entries[i].payload_ref.digest,
                               &refs[i].digest)) {
        amduat_asl_collection_refs_free(refs, i);
        amduat_asl_log_entries_free(entries, entries_len);
        return AMDUAT_ASL_COLLECTION_ERR_IO;
      }
    }
    *out_record_refs = refs;
    *out_len = entries_len;
  }

  amduat_asl_log_entries_free(entries, entries_len);
  return AMDUAT_ASL_COLLECTION_OK;
}

void amduat_asl_collection_refs_free(amduat_reference_t *refs,
                                     size_t refs_len) {
  if (refs == NULL) {
    return;
  }
  for (size_t i = 0u; i < refs_len; ++i) {
    amduat_reference_free(&refs[i]);
  }
  free(refs);
}