#include "amduat/asl/index_snapshot.h"

#include "amduat/asl/io.h"
#include "amduat/hash/asl1.h"

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

enum {
  AMDUAT_ASL_SNAPSHOT_MANIFEST_SEGMENT_SIZE = 40
};

static const uint8_t k_amduat_asl_snapshot_manifest_magic[
    AMDUAT_ASL_SNAPSHOT_MANIFEST_MAGIC_LEN] = {'A', 'S', 'L', 'S',
                                               'N', 'A', 'P', '1'};

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

static void amduat_asl_snapshot_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_snapshot_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 uint16_t amduat_asl_snapshot_load_u16_le(const uint8_t *data) {
  return (uint16_t)data[0] | ((uint16_t)data[1] << 8);
}

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

static uint64_t amduat_asl_snapshot_load_u64_le(const uint8_t *data) {
  return (uint64_t)data[0] | ((uint64_t)data[1] << 8) |
         ((uint64_t)data[2] << 16) | ((uint64_t)data[3] << 24) |
         ((uint64_t)data[4] << 32) | ((uint64_t)data[5] << 40) |
         ((uint64_t)data[6] << 48) | ((uint64_t)data[7] << 56);
}

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

static int amduat_asl_snapshot_compare_segments(const void *lhs,
                                                const void *rhs) {
  const amduat_asl_segment_seal_t *left =
      *(const amduat_asl_segment_seal_t *const *)lhs;
  const amduat_asl_segment_seal_t *right =
      *(const amduat_asl_segment_seal_t *const *)rhs;

  if (left->segment_id < right->segment_id) {
    return -1;
  }
  if (left->segment_id > right->segment_id) {
    return 1;
  }
  return memcmp(left->segment_hash, right->segment_hash,
                sizeof(left->segment_hash));
}

static int amduat_asl_snapshot_compare_tombstones(const void *lhs,
                                                  const void *rhs) {
  const amduat_asl_tombstone_entry_t *left =
      *(const amduat_asl_tombstone_entry_t *const *)lhs;
  const amduat_asl_tombstone_entry_t *right =
      *(const amduat_asl_tombstone_entry_t *const *)rhs;
  size_t min_len;
  int cmp;

  if (left->ref.hash_id < right->ref.hash_id) {
    return -1;
  }
  if (left->ref.hash_id > right->ref.hash_id) {
    return 1;
  }
  if (left->ref.digest.len < right->ref.digest.len) {
    return -1;
  }
  if (left->ref.digest.len > right->ref.digest.len) {
    return 1;
  }
  min_len = left->ref.digest.len;
  if (min_len != 0u) {
    cmp = memcmp(left->ref.digest.data, right->ref.digest.data, min_len);
    if (cmp != 0) {
      return cmp;
    }
  }
  if (left->tombstone_logseq < right->tombstone_logseq) {
    return -1;
  }
  if (left->tombstone_logseq > right->tombstone_logseq) {
    return 1;
  }
  return 0;
}

static bool amduat_asl_snapshot_hash_manifest(const uint8_t *bytes,
                                              size_t len,
                                              uint8_t out_hash[32]) {
  if (out_hash == NULL) {
    return true;
  }
  return amduat_hash_asl1_digest(AMDUAT_HASH_ASL1_ID_SHA256,
                                 amduat_octets(bytes, len), out_hash, 32);
}

void amduat_asl_snapshot_manifest_free(
    amduat_asl_snapshot_manifest_t *manifest) {
  size_t i;

  if (manifest == NULL) {
    return;
  }
  free(manifest->segments);
  manifest->segments = NULL;
  manifest->segments_len = 0;
  if (manifest->tombstones != NULL) {
    for (i = 0; i < manifest->tombstones_len; ++i) {
      amduat_reference_free(&manifest->tombstones[i].ref);
    }
  }
  free(manifest->tombstones);
  manifest->tombstones = NULL;
  manifest->tombstones_len = 0;
}

bool amduat_asl_snapshot_manifest_write(
    const char *path,
    const amduat_asl_snapshot_manifest_t *manifest,
    uint8_t out_hash[32]) {
  size_t i;
  size_t total_len;
  size_t offset;
  uint8_t *buffer;
  amduat_asl_segment_seal_t **segment_order;
  amduat_asl_tombstone_entry_t **tombstone_order;

  if (path == NULL || manifest == NULL) {
    return false;
  }
  if (manifest->segments_len != 0u && manifest->segments == NULL) {
    return false;
  }
  if (manifest->tombstones_len != 0u && manifest->tombstones == NULL) {
    return false;
  }

  if (manifest->segments_len > SIZE_MAX / AMDUAT_ASL_SNAPSHOT_MANIFEST_SEGMENT_SIZE) {
    return false;
  }

  total_len = 0;
  if (!amduat_asl_snapshot_add_size(&total_len,
                                    AMDUAT_ASL_SNAPSHOT_MANIFEST_HEADER_SIZE)) {
    return false;
  }
  if (!amduat_asl_snapshot_add_size(
          &total_len,
          manifest->segments_len *
              AMDUAT_ASL_SNAPSHOT_MANIFEST_SEGMENT_SIZE)) {
    return false;
  }

  for (i = 0; i < manifest->tombstones_len; ++i) {
    const amduat_asl_tombstone_entry_t *entry = &manifest->tombstones[i];
    if (entry->ref.digest.len == 0 ||
        (entry->ref.digest.len != 0u && entry->ref.digest.data == NULL)) {
      return false;
    }
    if (entry->ref.digest.len > UINT16_MAX) {
      return false;
    }
    if (!amduat_asl_snapshot_add_size(&total_len,
                                      16u + entry->ref.digest.len)) {
      return false;
    }
  }

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

  segment_order = NULL;
  tombstone_order = NULL;
  if (manifest->segments_len != 0u) {
    segment_order = (amduat_asl_segment_seal_t **)calloc(
        manifest->segments_len, sizeof(*segment_order));
    if (segment_order == NULL) {
      free(buffer);
      return false;
    }
    for (i = 0; i < manifest->segments_len; ++i) {
      segment_order[i] = &manifest->segments[i];
    }
    qsort(segment_order, manifest->segments_len, sizeof(*segment_order),
          amduat_asl_snapshot_compare_segments);
  }

  if (manifest->tombstones_len != 0u) {
    tombstone_order = (amduat_asl_tombstone_entry_t **)calloc(
        manifest->tombstones_len, sizeof(*tombstone_order));
    if (tombstone_order == NULL) {
      free(segment_order);
      free(buffer);
      return false;
    }
    for (i = 0; i < manifest->tombstones_len; ++i) {
      tombstone_order[i] = &manifest->tombstones[i];
    }
    qsort(tombstone_order, manifest->tombstones_len,
          sizeof(*tombstone_order), amduat_asl_snapshot_compare_tombstones);
  }

  offset = 0;
  memcpy(buffer + offset, k_amduat_asl_snapshot_manifest_magic,
         AMDUAT_ASL_SNAPSHOT_MANIFEST_MAGIC_LEN);
  offset += AMDUAT_ASL_SNAPSHOT_MANIFEST_MAGIC_LEN;
  amduat_asl_snapshot_store_u16_le(buffer + offset,
                                   AMDUAT_ASL_SNAPSHOT_MANIFEST_VERSION);
  offset += 2;
  amduat_asl_snapshot_store_u16_le(buffer + offset, 0);
  offset += 2;
  amduat_asl_snapshot_store_u32_le(buffer + offset,
                                   AMDUAT_ASL_SNAPSHOT_MANIFEST_HEADER_SIZE);
  offset += 4;
  amduat_asl_snapshot_store_u64_le(buffer + offset, manifest->snapshot_id);
  offset += 8;
  amduat_asl_snapshot_store_u64_le(buffer + offset, manifest->anchor_logseq);
  offset += 8;
  amduat_asl_snapshot_store_u64_le(buffer + offset, manifest->segments_len);
  offset += 8;
  amduat_asl_snapshot_store_u64_le(buffer + offset, manifest->tombstones_len);
  offset += 8;
  amduat_asl_snapshot_store_u32_le(buffer + offset,
                                   manifest->config.encoding_profile_id);
  offset += 4;
  amduat_asl_snapshot_store_u32_le(buffer + offset,
                                   manifest->config.hash_id);
  offset += 4;
  memset(buffer + offset, 0,
         AMDUAT_ASL_SNAPSHOT_MANIFEST_HEADER_SIZE - offset);
  offset = AMDUAT_ASL_SNAPSHOT_MANIFEST_HEADER_SIZE;

  for (i = 0; i < manifest->segments_len; ++i) {
    const amduat_asl_segment_seal_t *entry = segment_order[i];
    amduat_asl_snapshot_store_u64_le(buffer + offset, entry->segment_id);
    offset += 8;
    memcpy(buffer + offset, entry->segment_hash, sizeof(entry->segment_hash));
    offset += sizeof(entry->segment_hash);
  }

  for (i = 0; i < manifest->tombstones_len; ++i) {
    const amduat_asl_tombstone_entry_t *entry = tombstone_order[i];
    amduat_asl_snapshot_store_u32_le(buffer + offset, entry->ref.hash_id);
    offset += 4;
    amduat_asl_snapshot_store_u16_le(buffer + offset,
                                     (uint16_t)entry->ref.digest.len);
    offset += 2;
    amduat_asl_snapshot_store_u16_le(buffer + offset, 0);
    offset += 2;
    memcpy(buffer + offset, entry->ref.digest.data, entry->ref.digest.len);
    offset += entry->ref.digest.len;
    amduat_asl_snapshot_store_u64_le(buffer + offset,
                                     entry->tombstone_logseq);
    offset += 8;
  }

  free(segment_order);
  free(tombstone_order);

  if (offset != total_len) {
    free(buffer);
    return false;
  }

  if (!amduat_asl_snapshot_hash_manifest(buffer, total_len, out_hash)) {
    free(buffer);
    return false;
  }

  if (!amduat_asl_write_path(path, buffer, total_len)) {
    free(buffer);
    return false;
  }

  free(buffer);
  return true;
}

bool amduat_asl_snapshot_manifest_read(
    const char *path,
    amduat_asl_snapshot_manifest_t *out_manifest,
    uint8_t out_hash[32]) {
  uint8_t *buffer;
  size_t len;
  size_t offset;
  uint16_t version;
  uint16_t reserved0;
  uint32_t header_size;
  uint64_t segment_count;
  uint64_t tombstone_count;
  size_t i;

  if (path == NULL || out_manifest == NULL) {
    return false;
  }
  memset(out_manifest, 0, sizeof(*out_manifest));

  buffer = NULL;
  len = 0u;
  if (!amduat_asl_read_path(path, &buffer, &len)) {
    return false;
  }
  if (len < AMDUAT_ASL_SNAPSHOT_MANIFEST_HEADER_SIZE) {
    free(buffer);
    return false;
  }

  if (!amduat_asl_snapshot_hash_manifest(buffer, len, out_hash)) {
    free(buffer);
    return false;
  }

  if (memcmp(buffer, k_amduat_asl_snapshot_manifest_magic,
             AMDUAT_ASL_SNAPSHOT_MANIFEST_MAGIC_LEN) != 0) {
    free(buffer);
    return false;
  }

  offset = AMDUAT_ASL_SNAPSHOT_MANIFEST_MAGIC_LEN;
  version = amduat_asl_snapshot_load_u16_le(buffer + offset);
  offset += 2;
  reserved0 = amduat_asl_snapshot_load_u16_le(buffer + offset);
  offset += 2;
  header_size = amduat_asl_snapshot_load_u32_le(buffer + offset);
  offset += 4;

  if (version != AMDUAT_ASL_SNAPSHOT_MANIFEST_VERSION || reserved0 != 0 ||
      header_size != AMDUAT_ASL_SNAPSHOT_MANIFEST_HEADER_SIZE) {
    free(buffer);
    return false;
  }

  out_manifest->snapshot_id = amduat_asl_snapshot_load_u64_le(buffer + offset);
  offset += 8;
  out_manifest->anchor_logseq =
      amduat_asl_snapshot_load_u64_le(buffer + offset);
  offset += 8;
  segment_count = amduat_asl_snapshot_load_u64_le(buffer + offset);
  offset += 8;
  tombstone_count = amduat_asl_snapshot_load_u64_le(buffer + offset);
  offset += 8;
  out_manifest->config.encoding_profile_id =
      amduat_asl_snapshot_load_u32_le(buffer + offset);
  offset += 4;
  out_manifest->config.hash_id =
      amduat_asl_snapshot_load_u32_le(buffer + offset);
  offset += 4;
  offset = AMDUAT_ASL_SNAPSHOT_MANIFEST_HEADER_SIZE;

  if (segment_count > SIZE_MAX || tombstone_count > SIZE_MAX) {
    free(buffer);
    return false;
  }

  out_manifest->segments_len = (size_t)segment_count;
  if (out_manifest->segments_len != 0u) {
    if (out_manifest->segments_len >
        SIZE_MAX / AMDUAT_ASL_SNAPSHOT_MANIFEST_SEGMENT_SIZE) {
      free(buffer);
      return false;
    }
    out_manifest->segments = (amduat_asl_segment_seal_t *)calloc(
        out_manifest->segments_len, sizeof(*out_manifest->segments));
    if (out_manifest->segments == NULL) {
      free(buffer);
      return false;
    }
  }

  for (i = 0; i < out_manifest->segments_len; ++i) {
    amduat_asl_segment_seal_t *entry = &out_manifest->segments[i];
    if (len - offset < AMDUAT_ASL_SNAPSHOT_MANIFEST_SEGMENT_SIZE) {
      amduat_asl_snapshot_manifest_free(out_manifest);
      free(buffer);
      return false;
    }
    entry->segment_id = amduat_asl_snapshot_load_u64_le(buffer + offset);
    offset += 8;
    memcpy(entry->segment_hash, buffer + offset, sizeof(entry->segment_hash));
    offset += sizeof(entry->segment_hash);
  }

  out_manifest->tombstones_len = (size_t)tombstone_count;
  if (out_manifest->tombstones_len != 0u) {
    out_manifest->tombstones = (amduat_asl_tombstone_entry_t *)calloc(
        out_manifest->tombstones_len, sizeof(*out_manifest->tombstones));
    if (out_manifest->tombstones == NULL) {
      amduat_asl_snapshot_manifest_free(out_manifest);
      free(buffer);
      return false;
    }
  }

  for (i = 0; i < out_manifest->tombstones_len; ++i) {
    amduat_asl_tombstone_entry_t *entry = &out_manifest->tombstones[i];
    uint32_t hash_id;
    uint16_t digest_len;
    uint16_t tombstone_reserved;
    amduat_reference_t ref;
    uint64_t tombstone_logseq;

    if (len - offset < 16) {
      amduat_asl_snapshot_manifest_free(out_manifest);
      free(buffer);
      return false;
    }

    hash_id = amduat_asl_snapshot_load_u32_le(buffer + offset);
    offset += 4;
    if (hash_id > UINT16_MAX) {
      amduat_asl_snapshot_manifest_free(out_manifest);
      free(buffer);
      return false;
    }
    digest_len = amduat_asl_snapshot_load_u16_le(buffer + offset);
    offset += 2;
    tombstone_reserved = amduat_asl_snapshot_load_u16_le(buffer + offset);
    offset += 2;
    if (tombstone_reserved != 0 || digest_len == 0) {
      amduat_asl_snapshot_manifest_free(out_manifest);
      free(buffer);
      return false;
    }
    if (len - offset < (size_t)digest_len + 8u) {
      amduat_asl_snapshot_manifest_free(out_manifest);
      free(buffer);
      return false;
    }
    ref = amduat_reference((amduat_hash_id_t)hash_id,
                           amduat_octets(buffer + offset, digest_len));
    offset += digest_len;
    tombstone_logseq = amduat_asl_snapshot_load_u64_le(buffer + offset);
    offset += 8;
    if (!amduat_reference_clone(ref, &entry->ref)) {
      amduat_asl_snapshot_manifest_free(out_manifest);
      free(buffer);
      return false;
    }
    entry->tombstone_logseq = tombstone_logseq;
  }

  if (offset != len) {
    amduat_asl_snapshot_manifest_free(out_manifest);
    free(buffer);
    return false;
  }

  free(buffer);
  return true;
}