#include "amduat/enc/tgk1_edge.h"

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

static const uint8_t k_expected_edge_bytes[] = {
    0x00, 0x01, 0x00, 0x00, 0x00, 0x2a, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00,
    0x00, 0x22, 0x00, 0x01, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10,
    0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10,
    0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10,
    0x00, 0x00, 0x00, 0x22, 0x00, 0x01, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11,
    0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11,
    0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11,
    0x11, 0x11, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x22, 0x00, 0x01,
    0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20,
    0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20,
    0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x00, 0x00, 0x00, 0x22,
    0x00, 0x01, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30,
    0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30,
    0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30,
};

static const uint8_t k_empty_edge_bytes[] = {
    0x00, 0x01, 0x00, 0x00, 0x00, 0x2a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x22, 0x00, 0x01, 0x30, 0x30, 0x30, 0x30,
    0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30,
    0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30,
    0x30, 0x30, 0x30, 0x30,
};

static void fill_digest(uint8_t *out, uint8_t value) {
  memset(out, value, 32);
}

static amduat_reference_t make_ref(uint8_t value, uint8_t *storage) {
  fill_digest(storage, value);
  return amduat_reference(0x0001, amduat_octets(storage, 32));
}

static bool bytes_equal(amduat_octets_t bytes,
                        const uint8_t *expected,
                        size_t expected_len) {
  if (bytes.len != expected_len) {
    return false;
  }
  if (bytes.len == 0) {
    return true;
  }
  return memcmp(bytes.data, expected, expected_len) == 0;
}

static int test_edge_round_trip(void) {
  amduat_tgk_edge_body_t edge;
  amduat_reference_t from_refs[2];
  amduat_reference_t to_refs[1];
  amduat_octets_t encoded;
  amduat_tgk_edge_body_t decoded;
  uint8_t f0[32], f1[32], t0[32], p0[32];
  int exit_code = 1;

  memset(&edge, 0, sizeof(edge));
  edge.type = 0x2a;

  from_refs[0] = make_ref(0x10, f0);
  from_refs[1] = make_ref(0x11, f1);
  edge.from = from_refs;
  edge.from_len = 2;

  to_refs[0] = make_ref(0x20, t0);
  edge.to = to_refs;
  edge.to_len = 1;

  edge.payload = make_ref(0x30, p0);

  if (!amduat_enc_tgk1_edge_encode_v1(&edge, &encoded)) {
    fprintf(stderr, "encode failed\n");
    return exit_code;
  }

  if (!bytes_equal(encoded, k_expected_edge_bytes,
                   sizeof(k_expected_edge_bytes))) {
    fprintf(stderr, "encoded bytes mismatch\n");
    goto cleanup;
  }

  if (!amduat_enc_tgk1_edge_decode_v1(encoded, &decoded)) {
    fprintf(stderr, "decode failed\n");
    goto cleanup;
  }

  if (decoded.type != edge.type || decoded.from_len != 2 ||
      decoded.to_len != 1) {
    fprintf(stderr, "decoded field mismatch\n");
    goto cleanup_decoded;
  }

  if (!amduat_reference_eq(decoded.from[0], edge.from[0]) ||
      !amduat_reference_eq(decoded.from[1], edge.from[1]) ||
      !amduat_reference_eq(decoded.to[0], edge.to[0]) ||
      !amduat_reference_eq(decoded.payload, edge.payload)) {
    fprintf(stderr, "decoded references mismatch\n");
    goto cleanup_decoded;
  }

  exit_code = 0;

cleanup_decoded:
  amduat_enc_tgk1_edge_free(&decoded);
cleanup:
  free((void *)encoded.data);
  return exit_code;
}

static int test_invalid_edge_version(void) {
  uint8_t bad_bytes[sizeof(k_expected_edge_bytes)];
  amduat_octets_t bytes;
  amduat_tgk_edge_body_t decoded;

  memcpy(bad_bytes, k_expected_edge_bytes, sizeof(k_expected_edge_bytes));
  bad_bytes[1] = 0x02;

  bytes = amduat_octets(bad_bytes, sizeof(bad_bytes));
  if (amduat_enc_tgk1_edge_decode_v1(bytes, &decoded)) {
    fprintf(stderr, "invalid edge version accepted\n");
    amduat_enc_tgk1_edge_free(&decoded);
    return 1;
  }

  return 0;
}

static int test_empty_endpoints(void) {
  amduat_tgk_edge_body_t edge;
  amduat_octets_t encoded;
  amduat_tgk_edge_body_t decoded;
  uint8_t p0[32];

  memset(&edge, 0, sizeof(edge));
  edge.type = 0x2a;
  edge.payload = make_ref(0x30, p0);

  if (amduat_enc_tgk1_edge_encode_v1(&edge, &encoded)) {
    fprintf(stderr, "encode accepted empty endpoints\n");
    free((void *)encoded.data);
    return 1;
  }

  if (amduat_enc_tgk1_edge_decode_v1(
          amduat_octets(k_empty_edge_bytes, sizeof(k_empty_edge_bytes)),
          &decoded)) {
    fprintf(stderr, "decode accepted empty endpoints\n");
    amduat_enc_tgk1_edge_free(&decoded);
    return 1;
  }

  return 0;
}

static int test_trailing_bytes(void) {
  uint8_t bad_bytes[sizeof(k_expected_edge_bytes) + 1];
  amduat_octets_t bytes;
  amduat_tgk_edge_body_t decoded;

  memcpy(bad_bytes, k_expected_edge_bytes, sizeof(k_expected_edge_bytes));
  bad_bytes[sizeof(k_expected_edge_bytes)] = 0x00;

  bytes = amduat_octets(bad_bytes, sizeof(bad_bytes));
  if (amduat_enc_tgk1_edge_decode_v1(bytes, &decoded)) {
    fprintf(stderr, "trailing bytes accepted\n");
    amduat_enc_tgk1_edge_free(&decoded);
    return 1;
  }

  return 0;
}

int main(void) {
  if (test_edge_round_trip() != 0) {
    return 1;
  }
  if (test_invalid_edge_version() != 0) {
    return 1;
  }
  if (test_empty_endpoints() != 0) {
    return 1;
  }
  if (test_trailing_bytes() != 0) {
    return 1;
  }
  return 0;
}