#include "amduat/asl/collection_view.h"

#include "amduat/enc/asl1_core_codec.h"

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

enum {
  AMDUAT_ASL_COLLECTION_VIEW_MAGIC_LEN = 8,
  AMDUAT_ASL_COLLECTION_VIEW_VERSION = 1,
  AMDUAT_ASL_SNAPSHOT_INFO_MAGIC_LEN = 8,
  AMDUAT_ASL_SNAPSHOT_INFO_VERSION = 1,
  AMDUAT_ASL_LOG_RANGE_MAGIC_LEN = 8,
  AMDUAT_ASL_LOG_RANGE_VERSION = 1
};

static const uint8_t k_amduat_asl_collection_view_magic[
    AMDUAT_ASL_COLLECTION_VIEW_MAGIC_LEN] = {
        'A', 'S', 'L', 'C', 'V', 'W', '1', '\0'
};

static const uint8_t k_amduat_asl_snapshot_info_magic[
    AMDUAT_ASL_SNAPSHOT_INFO_MAGIC_LEN] = {
        'A', 'S', 'L', 'S', 'N', 'P', '1', '\0'
};

static const uint8_t k_amduat_asl_log_range_magic[
    AMDUAT_ASL_LOG_RANGE_MAGIC_LEN] = {
        'A', 'S', 'L', 'R', 'N', 'G', '1', '\0'
};

static void amduat_asl_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_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_read_u32_le(const uint8_t *data,
                                   size_t len,
                                   size_t *offset,
                                   uint32_t *out) {
  if (len - *offset < 4u) {
    return false;
  }
  *out = (uint32_t)data[*offset] |
         ((uint32_t)data[*offset + 1u] << 8) |
         ((uint32_t)data[*offset + 2u] << 16) |
         ((uint32_t)data[*offset + 3u] << 24);
  *offset += 4u;
  return true;
}

static bool amduat_asl_read_u64_le(const uint8_t *data,
                                   size_t len,
                                   size_t *offset,
                                   uint64_t *out) {
  if (len - *offset < 8u) {
    return false;
  }
  *out = (uint64_t)data[*offset] |
         ((uint64_t)data[*offset + 1u] << 8) |
         ((uint64_t)data[*offset + 2u] << 16) |
         ((uint64_t)data[*offset + 3u] << 24) |
         ((uint64_t)data[*offset + 4u] << 32) |
         ((uint64_t)data[*offset + 5u] << 40) |
         ((uint64_t)data[*offset + 6u] << 48) |
         ((uint64_t)data[*offset + 7u] << 56);
  *offset += 8u;
  return true;
}

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

static bool amduat_asl_encode_refs(const amduat_reference_t *refs,
                                   size_t refs_len,
                                   size_t *out_total_len) {
  size_t total = 0u;

  if (refs_len > UINT32_MAX) {
    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_add_size(&total, 4u + ref_bytes.len)) {
      free((void *)ref_bytes.data);
      return false;
    }
    free((void *)ref_bytes.data);
  }
  *out_total_len = total;
  return true;
}

static bool amduat_asl_write_refs(uint8_t *buffer,
                                  size_t total_len,
                                  size_t *offset,
                                  const amduat_reference_t *refs,
                                  size_t refs_len) {
  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 (total_len - *offset < 4u + ref_bytes.len) {
      free((void *)ref_bytes.data);
      return false;
    }
    amduat_asl_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);
  }
  return true;
}

static bool amduat_asl_decode_refs(amduat_octets_t bytes,
                                   size_t *offset,
                                   uint32_t refs_len,
                                   amduat_reference_t **out_refs,
                                   size_t *out_refs_len) {
  amduat_reference_t *refs = NULL;

  if (out_refs == NULL || out_refs_len == NULL) {
    return false;
  }
  *out_refs = NULL;
  *out_refs_len = 0u;

  if (refs_len == 0u) {
    return true;
  }
  if (refs_len > SIZE_MAX / sizeof(*refs)) {
    return false;
  }
  refs = (amduat_reference_t *)calloc(refs_len, sizeof(*refs));
  if (refs == NULL) {
    return false;
  }

  for (uint32_t i = 0u; i < refs_len; ++i) {
    uint32_t ref_len = 0u;
    amduat_octets_t ref_bytes;
    if (!amduat_asl_read_u32_le(bytes.data, bytes.len, offset, &ref_len)) {
      goto fail;
    }
    if (bytes.len - *offset < ref_len) {
      goto fail;
    }
    ref_bytes = amduat_octets(bytes.data + *offset, ref_len);
    if (!amduat_enc_asl1_core_decode_reference_v1(ref_bytes, &refs[i])) {
      goto fail;
    }
    *offset += ref_len;
  }

  *out_refs = refs;
  *out_refs_len = refs_len;
  return true;

fail:
  for (uint32_t i = 0u; i < refs_len; ++i) {
    amduat_reference_free(&refs[i]);
  }
  free(refs);
  return false;
}

bool amduat_asl_collection_snapshot_info_encode_v1(
    const amduat_asl_collection_snapshot_info_t *info,
    amduat_octets_t *out_bytes) {
  uint8_t *buffer;
  size_t total_len = 0u;
  size_t refs_len_bytes = 0u;
  size_t offset = 0u;

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

  if (!amduat_asl_add_size(&total_len,
                           AMDUAT_ASL_SNAPSHOT_INFO_MAGIC_LEN + 4u + 8u + 4u)) {
    return false;
  }
  if (!amduat_asl_encode_refs(info->refs, info->refs_len, &refs_len_bytes)) {
    return false;
  }
  if (!amduat_asl_add_size(&total_len, refs_len_bytes)) {
    return false;
  }

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

  memcpy(buffer + offset, k_amduat_asl_snapshot_info_magic,
         AMDUAT_ASL_SNAPSHOT_INFO_MAGIC_LEN);
  offset += AMDUAT_ASL_SNAPSHOT_INFO_MAGIC_LEN;
  amduat_asl_store_u32_le(buffer + offset, AMDUAT_ASL_SNAPSHOT_INFO_VERSION);
  offset += 4u;
  amduat_asl_store_u64_le(buffer + offset, info->snapshot_at_offset);
  offset += 8u;
  amduat_asl_store_u32_le(buffer + offset, (uint32_t)info->refs_len);
  offset += 4u;
  if (!amduat_asl_write_refs(buffer, total_len, &offset,
                             info->refs, info->refs_len)) {
    free(buffer);
    return false;
  }

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

bool amduat_asl_collection_snapshot_info_decode_v1(
    amduat_octets_t bytes,
    amduat_asl_collection_snapshot_info_t *out_info) {
  size_t offset = 0u;
  uint32_t version = 0u;
  uint32_t refs_len = 0u;

  if (out_info == NULL) {
    return false;
  }
  out_info->snapshot_at_offset = 0u;
  out_info->refs = NULL;
  out_info->refs_len = 0u;

  if (bytes.len < AMDUAT_ASL_SNAPSHOT_INFO_MAGIC_LEN + 4u + 8u + 4u ||
      bytes.data == NULL) {
    return false;
  }
  if (memcmp(bytes.data, k_amduat_asl_snapshot_info_magic,
             AMDUAT_ASL_SNAPSHOT_INFO_MAGIC_LEN) != 0) {
    return false;
  }
  offset += AMDUAT_ASL_SNAPSHOT_INFO_MAGIC_LEN;
  if (!amduat_asl_read_u32_le(bytes.data, bytes.len, &offset, &version) ||
      version != AMDUAT_ASL_SNAPSHOT_INFO_VERSION) {
    return false;
  }
  if (!amduat_asl_read_u64_le(bytes.data, bytes.len, &offset,
                              &out_info->snapshot_at_offset)) {
    return false;
  }
  if (!amduat_asl_read_u32_le(bytes.data, bytes.len, &offset, &refs_len)) {
    return false;
  }
  if (!amduat_asl_decode_refs(bytes, &offset, refs_len,
                              &out_info->refs, &out_info->refs_len)) {
    return false;
  }
  return offset == bytes.len;
}

void amduat_asl_collection_snapshot_info_free(
    amduat_asl_collection_snapshot_info_t *info) {
  if (info == NULL || info->refs == NULL) {
    return;
  }
  for (size_t i = 0u; i < info->refs_len; ++i) {
    amduat_reference_free(&info->refs[i]);
  }
  free(info->refs);
  info->refs = NULL;
  info->refs_len = 0u;
}

bool amduat_asl_log_range_encode_v1(const amduat_asl_log_range_t *range,
                                    amduat_octets_t *out_bytes) {
  uint8_t *buffer;
  size_t total_len = 0u;
  size_t refs_len_bytes = 0u;
  size_t offset = 0u;

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

  if (!amduat_asl_add_size(&total_len,
                           AMDUAT_ASL_LOG_RANGE_MAGIC_LEN + 4u + 8u + 8u + 4u)) {
    return false;
  }
  if (!amduat_asl_encode_refs(range->refs, range->refs_len, &refs_len_bytes)) {
    return false;
  }
  if (!amduat_asl_add_size(&total_len, refs_len_bytes)) {
    return false;
  }

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

  memcpy(buffer + offset, k_amduat_asl_log_range_magic,
         AMDUAT_ASL_LOG_RANGE_MAGIC_LEN);
  offset += AMDUAT_ASL_LOG_RANGE_MAGIC_LEN;
  amduat_asl_store_u32_le(buffer + offset, AMDUAT_ASL_LOG_RANGE_VERSION);
  offset += 4u;
  amduat_asl_store_u64_le(buffer + offset, range->start_offset);
  offset += 8u;
  amduat_asl_store_u64_le(buffer + offset, range->next_offset);
  offset += 8u;
  amduat_asl_store_u32_le(buffer + offset, (uint32_t)range->refs_len);
  offset += 4u;
  if (!amduat_asl_write_refs(buffer, total_len, &offset,
                             range->refs, range->refs_len)) {
    free(buffer);
    return false;
  }

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

bool amduat_asl_log_range_decode_v1(amduat_octets_t bytes,
                                    amduat_asl_log_range_t *out_range) {
  size_t offset = 0u;
  uint32_t version = 0u;
  uint32_t refs_len = 0u;

  if (out_range == NULL) {
    return false;
  }
  out_range->start_offset = 0u;
  out_range->next_offset = 0u;
  out_range->refs = NULL;
  out_range->refs_len = 0u;

  if (bytes.len <
          AMDUAT_ASL_LOG_RANGE_MAGIC_LEN + 4u + 8u + 8u + 4u ||
      bytes.data == NULL) {
    return false;
  }
  if (memcmp(bytes.data, k_amduat_asl_log_range_magic,
             AMDUAT_ASL_LOG_RANGE_MAGIC_LEN) != 0) {
    return false;
  }
  offset += AMDUAT_ASL_LOG_RANGE_MAGIC_LEN;
  if (!amduat_asl_read_u32_le(bytes.data, bytes.len, &offset, &version) ||
      version != AMDUAT_ASL_LOG_RANGE_VERSION) {
    return false;
  }
  if (!amduat_asl_read_u64_le(bytes.data, bytes.len, &offset,
                              &out_range->start_offset) ||
      !amduat_asl_read_u64_le(bytes.data, bytes.len, &offset,
                              &out_range->next_offset)) {
    return false;
  }
  if (!amduat_asl_read_u32_le(bytes.data, bytes.len, &offset, &refs_len)) {
    return false;
  }
  if (!amduat_asl_decode_refs(bytes, &offset, refs_len,
                              &out_range->refs, &out_range->refs_len)) {
    return false;
  }
  return offset == bytes.len;
}

void amduat_asl_log_range_free(amduat_asl_log_range_t *range) {
  if (range == NULL || range->refs == NULL) {
    return;
  }
  for (size_t i = 0u; i < range->refs_len; ++i) {
    amduat_reference_free(&range->refs[i]);
  }
  free(range->refs);
  range->refs = NULL;
  range->refs_len = 0u;
}

bool amduat_asl_collection_view_encode_v1(
    const amduat_asl_collection_view_t *view,
    amduat_octets_t *out_bytes) {
  uint8_t *buffer;
  size_t total_len = 0u;
  size_t refs_len_bytes = 0u;
  size_t offset = 0u;

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

  if (!amduat_asl_add_size(&total_len,
                           AMDUAT_ASL_COLLECTION_VIEW_MAGIC_LEN + 4u +
                           8u + 8u + 4u)) {
    return false;
  }
  if (!amduat_asl_encode_refs(view->refs, view->refs_len, &refs_len_bytes)) {
    return false;
  }
  if (!amduat_asl_add_size(&total_len, refs_len_bytes)) {
    return false;
  }

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

  memcpy(buffer + offset, k_amduat_asl_collection_view_magic,
         AMDUAT_ASL_COLLECTION_VIEW_MAGIC_LEN);
  offset += AMDUAT_ASL_COLLECTION_VIEW_MAGIC_LEN;
  amduat_asl_store_u32_le(buffer + offset,
                          AMDUAT_ASL_COLLECTION_VIEW_VERSION);
  offset += 4u;
  amduat_asl_store_u64_le(buffer + offset, view->computed_from_offset);
  offset += 8u;
  amduat_asl_store_u64_le(buffer + offset, view->computed_up_to_offset);
  offset += 8u;
  amduat_asl_store_u32_le(buffer + offset, (uint32_t)view->refs_len);
  offset += 4u;
  if (!amduat_asl_write_refs(buffer, total_len, &offset,
                             view->refs, view->refs_len)) {
    free(buffer);
    return false;
  }

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

bool amduat_asl_collection_view_decode_v1(
    amduat_octets_t bytes,
    amduat_asl_collection_view_t *out_view) {
  size_t offset = 0u;
  uint32_t version = 0u;
  uint32_t refs_len = 0u;

  if (out_view == NULL) {
    return false;
  }
  out_view->computed_from_offset = 0u;
  out_view->computed_up_to_offset = 0u;
  out_view->refs = NULL;
  out_view->refs_len = 0u;

  if (bytes.len <
          AMDUAT_ASL_COLLECTION_VIEW_MAGIC_LEN + 4u + 8u + 8u + 4u ||
      bytes.data == NULL) {
    return false;
  }
  if (memcmp(bytes.data, k_amduat_asl_collection_view_magic,
             AMDUAT_ASL_COLLECTION_VIEW_MAGIC_LEN) != 0) {
    return false;
  }
  offset += AMDUAT_ASL_COLLECTION_VIEW_MAGIC_LEN;
  if (!amduat_asl_read_u32_le(bytes.data, bytes.len, &offset, &version) ||
      version != AMDUAT_ASL_COLLECTION_VIEW_VERSION) {
    return false;
  }
  if (!amduat_asl_read_u64_le(bytes.data, bytes.len, &offset,
                              &out_view->computed_from_offset) ||
      !amduat_asl_read_u64_le(bytes.data, bytes.len, &offset,
                              &out_view->computed_up_to_offset)) {
    return false;
  }
  if (!amduat_asl_read_u32_le(bytes.data, bytes.len, &offset, &refs_len)) {
    return false;
  }
  if (!amduat_asl_decode_refs(bytes, &offset, refs_len,
                              &out_view->refs, &out_view->refs_len)) {
    return false;
  }
  return offset == bytes.len;
}

void amduat_asl_collection_view_free(amduat_asl_collection_view_t *view) {
  if (view == NULL || view->refs == NULL) {
    return;
  }
  for (size_t i = 0u; i < view->refs_len; ++i) {
    amduat_reference_free(&view->refs[i]);
  }
  free(view->refs);
  view->refs = NULL;
  view->refs_len = 0u;
}