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

#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>

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

static void 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 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 size_t build_artifact_ref(uint8_t *out,
                                 size_t out_cap,
                                 amduat_hash_id_t hash_id,
                                 const uint8_t *digest,
                                 size_t digest_len) {
  size_t total = 4 + 2 + 2 + digest_len;
  if (out_cap < total) {
    return 0;
  }
  store_u32_le(out, (uint32_t)hash_id);
  store_u16_le(out + 4, (uint16_t)digest_len);
  store_u16_le(out + 6, 0);
  memcpy(out + 8, digest, digest_len);
  return total;
}

static size_t build_tombstone_payload(uint8_t *out,
                                      size_t out_cap,
                                      amduat_hash_id_t hash_id,
                                      const uint8_t *digest,
                                      size_t digest_len) {
  size_t offset = 0;
  size_t ref_len = build_artifact_ref(out,
                                      out_cap,
                                      hash_id,
                                      digest,
                                      digest_len);
  if (ref_len == 0 || out_cap < ref_len + 8) {
    return 0;
  }
  offset += ref_len;
  store_u32_le(out + offset, 0);
  offset += 4;
  store_u32_le(out + offset, 0);
  offset += 4;
  return offset;
}

static size_t build_tombstone_lift_payload(uint8_t *out,
                                           size_t out_cap,
                                           amduat_hash_id_t hash_id,
                                           const uint8_t *digest,
                                           size_t digest_len,
                                           uint64_t tombstone_logseq) {
  size_t offset = 0;
  size_t ref_len = build_artifact_ref(out,
                                      out_cap,
                                      hash_id,
                                      digest,
                                      digest_len);
  if (ref_len == 0 || out_cap < ref_len + 8) {
    return 0;
  }
  offset += ref_len;
  store_u64_le(out + offset, tombstone_logseq);
  offset += 8;
  return offset;
}

static size_t build_segment_seal_payload(uint8_t *out,
                                         size_t out_cap,
                                         uint64_t segment_id,
                                         const uint8_t hash[32]) {
  if (out_cap < 8 + 32) {
    return 0;
  }
  store_u64_le(out, segment_id);
  memcpy(out + 8, hash, 32);
  return 8 + 32;
}

static int test_tombstone_lift_cutoff(void) {
  amduat_asl_log_record_t records[2];
  amduat_asl_replay_state_t state;
  uint8_t digest[32];
  uint8_t tombstone_payload[4 + 2 + 2 + 32 + 4 + 4];
  uint8_t lift_payload[4 + 2 + 2 + 32 + 8];
  size_t tombstone_len;
  size_t lift_len;
  amduat_reference_t ref;

  memset(digest, 0x5a, sizeof(digest));
  tombstone_len = build_tombstone_payload(tombstone_payload,
                                          sizeof(tombstone_payload),
                                          AMDUAT_HASH_ASL1_ID_SHA256,
                                          digest,
                                          sizeof(digest));
  lift_len = build_tombstone_lift_payload(lift_payload,
                                          sizeof(lift_payload),
                                          AMDUAT_HASH_ASL1_ID_SHA256,
                                          digest,
                                          sizeof(digest),
                                          20);
  if (tombstone_len == 0 || lift_len == 0) {
    fprintf(stderr, "payload build failed\n");
    return 1;
  }

  memset(records, 0, sizeof(records));
  records[0].logseq = 20;
  records[0].record_type = AMDUAT_ASL_LOG_RECORD_TOMBSTONE;
  records[0].payload = amduat_octets(tombstone_payload, tombstone_len);
  records[1].logseq = 40;
  records[1].record_type = AMDUAT_ASL_LOG_RECORD_TOMBSTONE_LIFT;
  records[1].payload = amduat_octets(lift_payload, lift_len);

  if (!amduat_asl_replay_init(&state)) {
    fprintf(stderr, "replay init failed\n");
    return 1;
  }
  if (!amduat_asl_replay_apply_log(records, 2, 30, &state)) {
    fprintf(stderr, "replay apply failed\n");
    amduat_asl_replay_free(&state);
    return 1;
  }
  if (state.tombstones_len != 1) {
    fprintf(stderr, "tombstone not applied\n");
    amduat_asl_replay_free(&state);
    return 1;
  }
  ref = amduat_reference(AMDUAT_HASH_ASL1_ID_SHA256,
                         amduat_octets(digest, sizeof(digest)));
  if (state.tombstones[0].tombstone_logseq != 20 ||
      !amduat_reference_eq(state.tombstones[0].ref, ref)) {
    fprintf(stderr, "tombstone mismatch\n");
    amduat_asl_replay_free(&state);
    return 1;
  }
  amduat_asl_replay_free(&state);

  if (!amduat_asl_replay_init(&state)) {
    fprintf(stderr, "replay init failed\n");
    return 1;
  }
  if (!amduat_asl_replay_apply_log(records, 2, 40, &state)) {
    fprintf(stderr, "replay apply failed\n");
    amduat_asl_replay_free(&state);
    return 1;
  }
  if (state.tombstones_len != 0) {
    fprintf(stderr, "tombstone lift failed\n");
    amduat_asl_replay_free(&state);
    return 1;
  }
  amduat_asl_replay_free(&state);
  return 0;
}

static int test_unknown_record_skip(void) {
  amduat_asl_log_record_t records[2];
  amduat_asl_replay_state_t state;
  uint8_t unknown_payload[1] = {0x01};
  uint8_t seal_payload[8 + 32];
  uint8_t seal_hash[32];
  size_t seal_len;

  memset(seal_hash, 0xab, sizeof(seal_hash));
  seal_len = build_segment_seal_payload(seal_payload,
                                        sizeof(seal_payload),
                                        9,
                                        seal_hash);
  if (seal_len == 0) {
    fprintf(stderr, "seal payload build failed\n");
    return 1;
  }

  memset(records, 0, sizeof(records));
  records[0].logseq = 10;
  records[0].record_type = 0x99;
  records[0].payload = amduat_octets(unknown_payload, sizeof(unknown_payload));
  records[1].logseq = 11;
  records[1].record_type = AMDUAT_ASL_LOG_RECORD_SEGMENT_SEAL;
  records[1].payload = amduat_octets(seal_payload, seal_len);

  if (!amduat_asl_replay_init(&state)) {
    fprintf(stderr, "replay init failed\n");
    return 1;
  }
  if (!amduat_asl_replay_apply_log(records, 2, 11, &state)) {
    fprintf(stderr, "replay apply failed\n");
    amduat_asl_replay_free(&state);
    return 1;
  }
  if (state.segments_len != 1) {
    fprintf(stderr, "segment seal missing\n");
    amduat_asl_replay_free(&state);
    return 1;
  }
  amduat_asl_replay_free(&state);
  return 0;
}

static int test_multiple_seals_latest(void) {
  amduat_asl_log_record_t records[2];
  amduat_asl_replay_state_t state;
  uint8_t seal_payload_a[8 + 32];
  uint8_t seal_payload_b[8 + 32];
  uint8_t hash_a[32];
  uint8_t hash_b[32];
  size_t seal_len_a;
  size_t seal_len_b;

  memset(hash_a, 0x11, sizeof(hash_a));
  memset(hash_b, 0x22, sizeof(hash_b));
  seal_len_a = build_segment_seal_payload(seal_payload_a,
                                          sizeof(seal_payload_a),
                                          5,
                                          hash_a);
  seal_len_b = build_segment_seal_payload(seal_payload_b,
                                          sizeof(seal_payload_b),
                                          5,
                                          hash_b);
  if (seal_len_a == 0 || seal_len_b == 0) {
    fprintf(stderr, "seal payload build failed\n");
    return 1;
  }

  memset(records, 0, sizeof(records));
  records[0].logseq = 10;
  records[0].record_type = AMDUAT_ASL_LOG_RECORD_SEGMENT_SEAL;
  records[0].payload = amduat_octets(seal_payload_a, seal_len_a);
  records[1].logseq = 20;
  records[1].record_type = AMDUAT_ASL_LOG_RECORD_SEGMENT_SEAL;
  records[1].payload = amduat_octets(seal_payload_b, seal_len_b);

  if (!amduat_asl_replay_init(&state)) {
    fprintf(stderr, "replay init failed\n");
    return 1;
  }
  if (!amduat_asl_replay_apply_log(records, 2, 20, &state)) {
    fprintf(stderr, "replay apply failed\n");
    amduat_asl_replay_free(&state);
    return 1;
  }
  if (state.segments_len != 1) {
    fprintf(stderr, "segment seal count mismatch\n");
    amduat_asl_replay_free(&state);
    return 1;
  }
  if (memcmp(state.segments[0].segment_hash, hash_b, sizeof(hash_b)) != 0) {
    fprintf(stderr, "segment seal not updated\n");
    amduat_asl_replay_free(&state);
    return 1;
  }
  amduat_asl_replay_free(&state);
  return 0;
}

int main(void) {
  if (test_tombstone_lift_cutoff() != 0) {
    return 1;
  }
  if (test_unknown_record_skip() != 0) {
    return 1;
  }
  if (test_multiple_seals_latest() != 0) {
    return 1;
  }
  return 0;
}