#include "amduat/enc/fer1_receipt.h"
#include "amduat/fer/receipt.h"

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

static const uint8_t k_expected_receipt_bytes[] = {
    0x00, 0x01, 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, 0x22, 0x00, 0x01, 0x22, 0x22,
    0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22,
    0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22,
    0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x00, 0x00, 0x00, 0x22, 0x00, 0x01,
    0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33,
    0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33,
    0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x00, 0x00, 0x00, 0x06,
    0x74, 0x65, 0x73, 0x74, 0x65, 0x72, 0x00, 0x00, 0x00, 0x22, 0x00, 0x01,
    0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44,
    0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44,
    0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x00, 0x00, 0x00, 0x02,
    0x00, 0x00, 0x00, 0x22, 0x00, 0x01, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50,
    0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50,
    0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50,
    0x50, 0x50, 0x00, 0x00, 0x00, 0x22, 0x00, 0x01, 0x51, 0x51, 0x51, 0x51,
    0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51,
    0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51,
    0x51, 0x51, 0x51, 0x51, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0x22,
    0x00, 0x01, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50,
    0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50,
    0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x00, 0x00,
    0x00, 0x22, 0x00, 0x01, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44,
    0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44,
    0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44,
    0x01, 0x00, 0x00, 0x00, 0x22, 0x00, 0x01, 0x60, 0x60, 0x60, 0x60, 0x60,
    0x60, 0x60, 0x60, 0x60, 0x60, 0x60, 0x60, 0x60, 0x60, 0x60, 0x60, 0x60,
    0x60, 0x60, 0x60, 0x60, 0x60, 0x60, 0x60, 0x60, 0x60, 0x60, 0x60, 0x60,
    0x60, 0x60, 0x60, 0x00, 0x00, 0x00, 0x03, 0xaa, 0xbb, 0xcc, 0x00, 0x00,
    0x00, 0x22, 0x00, 0x01, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51,
    0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51,
    0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51, 0x51,
    0x00, 0x00, 0x00, 0x22, 0x00, 0x01, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44,
    0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44,
    0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44,
    0x44, 0x44, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x14
};

static const uint8_t k_expected_receipt_helper_bytes[] = {
    0x00, 0x01, 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, 0x22, 0x00, 0x01, 0x22, 0x22,
    0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22,
    0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22,
    0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x00, 0x00, 0x00, 0x22, 0x00, 0x01,
    0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33,
    0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33,
    0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x00, 0x00, 0x00, 0x06,
    0x74, 0x65, 0x73, 0x74, 0x65, 0x72, 0x00, 0x00, 0x00, 0x22, 0x00, 0x01,
    0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44,
    0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44,
    0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x00, 0x00, 0x00, 0x01,
    0x00, 0x00, 0x00, 0x22, 0x00, 0x01, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50,
    0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50,
    0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50,
    0x50, 0x50, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x22, 0x00, 0x01,
    0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50,
    0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50,
    0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x50, 0x00, 0x00, 0x00, 0x22,
    0x00, 0x01, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44,
    0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44,
    0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x44, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0a, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x14,
};

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_receipt_round_trip(void) {
  amduat_fer1_receipt_t receipt;
  amduat_octets_t encoded;
  amduat_fer1_receipt_t decoded;
  amduat_reference_t executor_refs[2];
  amduat_fer1_parity_entry_t parity[2];
  uint8_t f0[32], i0[32], e0[32], o0[32];
  uint8_t ex0[32], ex1[32], sb0[32];
  uint8_t digest0[] = {0xaa, 0xbb, 0xcc};
  int exit_code = 1;

  memset(&receipt, 0, sizeof(receipt));
  receipt.fer1_version = 1;
  receipt.function_ref = make_ref(0x11, f0);
  receipt.input_manifest_ref = make_ref(0x22, i0);
  receipt.environment_ref = make_ref(0x33, e0);
  receipt.evaluator_id = amduat_octets("tester", 6);
  receipt.output_ref = make_ref(0x44, o0);
  receipt.started_at = 10;
  receipt.completed_at = 20;

  executor_refs[0] = make_ref(0x50, ex0);
  executor_refs[1] = make_ref(0x51, ex1);
  receipt.executor_refs = executor_refs;
  receipt.executor_refs_len = 2;

  memset(parity, 0, sizeof(parity));
  parity[0].executor_ref = executor_refs[0];
  parity[0].output_ref = receipt.output_ref;
  parity[0].has_sbom_ref = true;
  parity[0].sbom_ref = make_ref(0x60, sb0);
  parity[0].parity_digest = amduat_octets(digest0, sizeof(digest0));

  parity[1].executor_ref = executor_refs[1];
  parity[1].output_ref = receipt.output_ref;
  parity[1].has_sbom_ref = false;
  parity[1].parity_digest = amduat_octets(NULL, 0);

  receipt.parity = parity;
  receipt.parity_len = 2;

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

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

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

  if (!amduat_reference_eq(decoded.function_ref, receipt.function_ref) ||
      !amduat_reference_eq(decoded.input_manifest_ref,
                           receipt.input_manifest_ref) ||
      !amduat_reference_eq(decoded.environment_ref,
                           receipt.environment_ref) ||
      !amduat_reference_eq(decoded.output_ref, receipt.output_ref)) {
    fprintf(stderr, "decoded refs mismatch\n");
    goto cleanup_decoded;
  }

  if (!amduat_octets_eq(decoded.evaluator_id, receipt.evaluator_id) ||
      decoded.executor_refs_len != 2 || decoded.parity_len != 2 ||
      decoded.started_at != receipt.started_at ||
      decoded.completed_at != receipt.completed_at) {
    fprintf(stderr, "decoded fields mismatch\n");
    goto cleanup_decoded;
  }

  if (!amduat_reference_eq(decoded.executor_refs[0], executor_refs[0]) ||
      !amduat_reference_eq(decoded.executor_refs[1], executor_refs[1]) ||
      !amduat_reference_eq(decoded.parity[0].executor_ref, executor_refs[0]) ||
      !amduat_reference_eq(decoded.parity[0].output_ref,
                           receipt.output_ref) ||
      !decoded.parity[0].has_sbom_ref ||
      !amduat_reference_eq(decoded.parity[0].sbom_ref, parity[0].sbom_ref) ||
      !amduat_octets_eq(decoded.parity[0].parity_digest,
                        parity[0].parity_digest) ||
      !amduat_reference_eq(decoded.parity[1].executor_ref, executor_refs[1]) ||
      !amduat_reference_eq(decoded.parity[1].output_ref,
                           receipt.output_ref) ||
      decoded.parity[1].has_sbom_ref ||
      decoded.parity[1].parity_digest.len != 0) {
    fprintf(stderr, "decoded parity mismatch\n");
    goto cleanup_decoded;
  }

  exit_code = 0;

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

static int test_invalid_receipt_version(void) {
  uint8_t bad_bytes[sizeof(k_expected_receipt_bytes)];
  amduat_octets_t bytes;
  amduat_fer1_receipt_t decoded;

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

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

  return 0;
}

static int test_receipt_helper(void) {
  amduat_pel_surface_execution_result_t pel_result;
  amduat_artifact_t artifact;
  amduat_reference_t output_ref;
  uint8_t f0[32], i0[32], e0[32], o0[32];
  uint8_t ex0[32];

  memset(&pel_result, 0, sizeof(pel_result));
  pel_result.pel1_version = 1;
  pel_result.program_ref = make_ref(0x11, f0);
  pel_result.output_refs = &output_ref;
  pel_result.output_refs_len = 1;
  output_ref = make_ref(0x44, o0);

  if (!amduat_fer1_receipt_from_pel_result(
          &pel_result,
          make_ref(0x22, i0),
          make_ref(0x33, e0),
          amduat_octets("tester", 6),
          make_ref(0x50, ex0),
          false,
          amduat_reference(0, amduat_octets(NULL, 0)),
          amduat_octets(NULL, 0),
          10,
          20,
          &artifact)) {
    fprintf(stderr, "helper failed\n");
    return 1;
  }

  if (!bytes_equal(artifact.bytes, k_expected_receipt_helper_bytes,
                   sizeof(k_expected_receipt_helper_bytes))) {
    fprintf(stderr, "helper bytes mismatch\n");
    amduat_artifact_free(&artifact);
    return 1;
  }

  amduat_artifact_free(&artifact);
  return 0;
}

int main(void) {
  if (test_receipt_round_trip() != 0) {
    return 1;
  }
  if (test_invalid_receipt_version() != 0) {
    return 1;
  }
  if (test_receipt_helper() != 0) {
    return 1;
  }
  return 0;
}