#include "amduat/asl/index_replay.h"

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

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

static bool amduat_asl_replay_read_u16_le(amduat_asl_replay_cursor_t *cur,
                                          uint16_t *out) {
  const uint8_t *data;

  if (cur->len - cur->offset < 2) {
    return false;
  }
  data = cur->data + cur->offset;
  *out = (uint16_t)data[0] | ((uint16_t)data[1] << 8);
  cur->offset += 2;
  return true;
}

static bool amduat_asl_replay_read_u32_le(amduat_asl_replay_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] | ((uint32_t)data[1] << 8) |
         ((uint32_t)data[2] << 16) | ((uint32_t)data[3] << 24);
  cur->offset += 4;
  return true;
}

static bool amduat_asl_replay_read_u64_le(amduat_asl_replay_cursor_t *cur,
                                          uint64_t *out) {
  const uint8_t *data;

  if (cur->len - cur->offset < 8) {
    return false;
  }
  data = cur->data + cur->offset;
  *out = (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);
  cur->offset += 8;
  return true;
}

static bool amduat_asl_replay_parse_artifact_ref(
    amduat_asl_replay_cursor_t *cur,
    amduat_reference_t *out_ref) {
  uint32_t hash_id_raw;
  uint16_t digest_len;
  uint16_t reserved0;
  const uint8_t *digest;

  if (!amduat_asl_replay_read_u32_le(cur, &hash_id_raw) ||
      !amduat_asl_replay_read_u16_le(cur, &digest_len) ||
      !amduat_asl_replay_read_u16_le(cur, &reserved0)) {
    return false;
  }
  if (hash_id_raw > UINT16_MAX || digest_len == 0 || reserved0 != 0) {
    return false;
  }
  if (cur->len - cur->offset < digest_len) {
    return false;
  }
  digest = cur->data + cur->offset;
  cur->offset += digest_len;
  *out_ref = amduat_reference((amduat_hash_id_t)hash_id_raw,
                              amduat_octets(digest, digest_len));
  return true;
}

static bool amduat_asl_replay_parse_segment_seal(
    amduat_octets_t payload,
    amduat_asl_segment_seal_t *out) {
  amduat_asl_replay_cursor_t cur;
  uint64_t segment_id;

  if (payload.len < 8 + 32 || payload.data == NULL || out == NULL) {
    return false;
  }
  cur.data = payload.data;
  cur.len = payload.len;
  cur.offset = 0;
  if (!amduat_asl_replay_read_u64_le(&cur, &segment_id)) {
    return false;
  }
  if (cur.len - cur.offset < 32) {
    return false;
  }
  out->segment_id = segment_id;
  memcpy(out->segment_hash, cur.data + cur.offset, 32);
  return true;
}

static bool amduat_asl_replay_parse_tombstone(
    amduat_octets_t payload,
    amduat_reference_t *out_ref) {
  amduat_asl_replay_cursor_t cur;
  uint32_t scope;
  uint32_t reason;

  if (payload.len == 0 || payload.data == NULL || out_ref == NULL) {
    return false;
  }
  cur.data = payload.data;
  cur.len = payload.len;
  cur.offset = 0;
  if (!amduat_asl_replay_parse_artifact_ref(&cur, out_ref)) {
    return false;
  }
  if (!amduat_asl_replay_read_u32_le(&cur, &scope) ||
      !amduat_asl_replay_read_u32_le(&cur, &reason)) {
    return false;
  }
  (void)scope;
  (void)reason;
  return true;
}

static bool amduat_asl_replay_parse_tombstone_lift(
    amduat_octets_t payload,
    amduat_reference_t *out_ref,
    uint64_t *out_logseq) {
  amduat_asl_replay_cursor_t cur;
  uint64_t tombstone_logseq;

  if (payload.len == 0 || payload.data == NULL || out_ref == NULL ||
      out_logseq == NULL) {
    return false;
  }
  cur.data = payload.data;
  cur.len = payload.len;
  cur.offset = 0;
  if (!amduat_asl_replay_parse_artifact_ref(&cur, out_ref) ||
      !amduat_asl_replay_read_u64_le(&cur, &tombstone_logseq)) {
    return false;
  }
  *out_logseq = tombstone_logseq;
  return true;
}

static bool amduat_asl_replay_update_segment(
    amduat_asl_replay_state_t *state,
    const amduat_asl_segment_seal_t *seal) {
  size_t i;
  amduat_asl_segment_seal_t *next;

  for (i = 0; i < state->segments_len; ++i) {
    if (state->segments[i].segment_id == seal->segment_id) {
      memcpy(state->segments[i].segment_hash, seal->segment_hash, 32);
      return true;
    }
  }

  next = (amduat_asl_segment_seal_t *)realloc(
      state->segments,
      (state->segments_len + 1u) * sizeof(*state->segments));
  if (next == NULL) {
    return false;
  }
  state->segments = next;
  state->segments[state->segments_len] = *seal;
  state->segments_len += 1u;
  return true;
}

static bool amduat_asl_replay_add_tombstone(
    amduat_asl_replay_state_t *state,
    amduat_reference_t ref,
    uint64_t logseq) {
  amduat_asl_tombstone_entry_t *next;
  amduat_reference_t stored;

  if (!amduat_reference_clone(ref, &stored)) {
    return false;
  }
  next = (amduat_asl_tombstone_entry_t *)realloc(
      state->tombstones,
      (state->tombstones_len + 1u) * sizeof(*state->tombstones));
  if (next == NULL) {
    amduat_reference_free(&stored);
    return false;
  }
  state->tombstones = next;
  state->tombstones[state->tombstones_len].ref = stored;
  state->tombstones[state->tombstones_len].tombstone_logseq = logseq;
  state->tombstones_len += 1u;
  return true;
}

static void amduat_asl_replay_remove_tombstone(
    amduat_asl_replay_state_t *state,
    amduat_reference_t ref,
    uint64_t tombstone_logseq) {
  size_t i;

  for (i = 0; i < state->tombstones_len; ++i) {
    if (state->tombstones[i].tombstone_logseq != tombstone_logseq) {
      continue;
    }
    if (!amduat_reference_eq(state->tombstones[i].ref, ref)) {
      continue;
    }
    amduat_reference_free(&state->tombstones[i].ref);
    if (i + 1u < state->tombstones_len) {
      state->tombstones[i] = state->tombstones[state->tombstones_len - 1u];
    }
    state->tombstones_len -= 1u;
    return;
  }
}

bool amduat_asl_replay_init(amduat_asl_replay_state_t *out) {
  if (out == NULL) {
    return false;
  }
  out->segments = NULL;
  out->segments_len = 0;
  out->tombstones = NULL;
  out->tombstones_len = 0;
  out->state.snapshot_id = 0;
  out->state.log_position = 0;
  return true;
}

void amduat_asl_replay_free(amduat_asl_replay_state_t *state) {
  size_t i;

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

bool amduat_asl_replay_apply_log(
    const amduat_asl_log_record_t *records,
    size_t record_count,
    amduat_asl_log_position_t log_position,
    amduat_asl_replay_state_t *state) {
  size_t i;

  if (state == NULL) {
    return false;
  }
  if (record_count != 0 && records == NULL) {
    return false;
  }

  for (i = 0; i < record_count; ++i) {
    const amduat_asl_log_record_t *record = &records[i];
    amduat_asl_segment_seal_t seal;
    amduat_reference_t ref;
    uint64_t tombstone_logseq;

    if (record->logseq > log_position) {
      break;
    }

    switch (record->record_type) {
      case AMDUAT_ASL_LOG_RECORD_SEGMENT_SEAL:
        if (!amduat_asl_replay_parse_segment_seal(record->payload, &seal)) {
          return false;
        }
        if (!amduat_asl_replay_update_segment(state, &seal)) {
          return false;
        }
        break;
      case AMDUAT_ASL_LOG_RECORD_TOMBSTONE:
        if (!amduat_asl_replay_parse_tombstone(record->payload, &ref)) {
          return false;
        }
        if (!amduat_asl_replay_add_tombstone(state, ref, record->logseq)) {
          return false;
        }
        break;
      case AMDUAT_ASL_LOG_RECORD_TOMBSTONE_LIFT:
        if (!amduat_asl_replay_parse_tombstone_lift(record->payload,
                                                    &ref,
                                                    &tombstone_logseq)) {
          return false;
        }
        amduat_asl_replay_remove_tombstone(state, ref, tombstone_logseq);
        break;
      default:
        break;
    }
  }

  state->state.log_position = log_position;
  return true;
}