#include "amduat/asl/index_accel.h"
#include "amduat/asl/index_bloom.h"

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

static int test_routing_key_layout(void) {
  uint8_t digest_bytes[2] = {0xaa, 0xbb};
  amduat_reference_t ref =
      amduat_reference(0x1234u, amduat_octets(digest_bytes, sizeof(digest_bytes)));
  amduat_octets_t key;
  int exit_code = 1;

  if (!amduat_asl_index_accel_routing_key_from_ref(
          ref, true, amduat_type_tag(0x11223344u), &key)) {
    fprintf(stderr, "routing key build failed\n");
    return exit_code;
  }

  if (key.len != 2u + 2u + sizeof(digest_bytes) + 1u + 4u) {
    fprintf(stderr, "routing key length mismatch\n");
    goto cleanup;
  }
  if (key.data[0] != 0x34u || key.data[1] != 0x12u ||
      key.data[2] != 0x02u || key.data[3] != 0x00u) {
    fprintf(stderr, "routing key header mismatch\n");
    goto cleanup;
  }
  if (key.data[4] != 0xaau || key.data[5] != 0xbbu) {
    fprintf(stderr, "routing key digest mismatch\n");
    goto cleanup;
  }
  if (key.data[6] != 0x01u) {
    fprintf(stderr, "routing key type tag flag mismatch\n");
    goto cleanup;
  }
  if (key.data[7] != 0x44u || key.data[8] != 0x33u ||
      key.data[9] != 0x22u || key.data[10] != 0x11u) {
    fprintf(stderr, "routing key type tag bytes mismatch\n");
    goto cleanup;
  }

  exit_code = 0;

cleanup:
  amduat_octets_free(&key);
  return exit_code;
}

static int test_routing_key_absence(void) {
  uint8_t digest_bytes[1] = {0x7f};
  amduat_reference_t ref =
      amduat_reference(0x0001u, amduat_octets(digest_bytes, sizeof(digest_bytes)));
  amduat_octets_t key;
  int exit_code = 1;

  if (!amduat_asl_index_accel_routing_key_from_ref(
          ref, false, amduat_type_tag(0xdeadbeefu), &key)) {
    fprintf(stderr, "routing key build failed\n");
    return exit_code;
  }

  if (key.data[4] != 0x7fu) {
    fprintf(stderr, "routing key digest mismatch\n");
    goto cleanup;
  }
  if (key.data[5] != 0x00u) {
    fprintf(stderr, "routing key absence flag mismatch\n");
    goto cleanup;
  }
  if (key.data[6] != 0x00u || key.data[7] != 0x00u ||
      key.data[8] != 0x00u || key.data[9] != 0x00u) {
    fprintf(stderr, "routing key absence tag mismatch\n");
    goto cleanup;
  }

  exit_code = 0;

cleanup:
  amduat_octets_free(&key);
  return exit_code;
}

static int test_shard_determinism(void) {
  uint8_t digest_bytes[3] = {0x01, 0x02, 0x03};
  amduat_reference_t ref =
      amduat_reference(0x00f0u, amduat_octets(digest_bytes, sizeof(digest_bytes)));
  uint16_t shard_a;
  uint16_t shard_b;

  shard_a = amduat_asl_index_accel_shard_for_ref(
      ref, false, amduat_type_tag(0u), 8u);
  shard_b = amduat_asl_index_accel_shard_for_ref(
      ref, false, amduat_type_tag(0u), 8u);

  if (shard_a != shard_b) {
    fprintf(stderr, "shard selection mismatch\n");
    return 1;
  }
  if (shard_a >= 8u) {
    fprintf(stderr, "shard out of range\n");
    return 1;
  }
  return 0;
}

static int test_bloom_advisory(void) {
  uint8_t digest_bytes[4] = {0x10, 0x20, 0x30, 0x40};
  amduat_octets_t bloom;
  amduat_octets_t empty = amduat_octets(NULL, 0u);
  int exit_code = 1;

  if (!amduat_asl_index_bloom_init(&bloom)) {
    fprintf(stderr, "bloom init failed\n");
    return exit_code;
  }
  if (!amduat_asl_index_bloom_add(
          bloom, 0x0001u, amduat_octets(digest_bytes, sizeof(digest_bytes)))) {
    fprintf(stderr, "bloom add failed\n");
    goto cleanup;
  }
  if (!amduat_asl_index_bloom_maybe_contains(
          bloom, 0x0001u, amduat_octets(digest_bytes, sizeof(digest_bytes)))) {
    fprintf(stderr, "bloom false negative\n");
    goto cleanup;
  }
  if (!amduat_asl_index_bloom_maybe_contains(
          empty, 0x0002u, amduat_octets(NULL, 0u))) {
    fprintf(stderr, "empty bloom should be advisory\n");
    goto cleanup;
  }

  exit_code = 0;

cleanup:
  amduat_octets_free(&bloom);
  return exit_code;
}

int main(void) {
  if (test_routing_key_layout() != 0) {
    return 1;
  }
  if (test_routing_key_absence() != 0) {
    return 1;
  }
  if (test_shard_determinism() != 0) {
    return 1;
  }
  return test_bloom_advisory();
}