#include "amduat/enc/pel1_result.h"

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

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 void init_base_result(amduat_pel_surface_execution_result_t *result,
                             uint8_t *scheme_digest,
                             uint8_t *program_digest) {
  memset(result, 0, sizeof(*result));
  result->pel1_version = 1;
  result->scheme_ref = make_ref(0x01, scheme_digest);
  result->program_ref = make_ref(0x02, program_digest);
  result->input_refs = NULL;
  result->input_refs_len = 0;
  result->output_refs = NULL;
  result->output_refs_len = 0;
  result->has_params_ref = false;
  result->has_store_failure = false;
  result->has_trace_ref = false;

  result->core_result.pel1_version = 1;
  result->core_result.status = AMDUAT_PEL_EXEC_STATUS_OK;
  result->core_result.scheme_ref = result->scheme_ref;
  result->core_result.summary.kind = AMDUAT_PEL_EXEC_ERROR_NONE;
  result->core_result.summary.status_code = 0;
  result->core_result.diagnostics = NULL;
  result->core_result.diagnostics_len = 0;
}

static int test_encode_status_invariant(void) {
  amduat_pel_surface_execution_result_t result;
  amduat_octets_t encoded = amduat_octets(NULL, 0);
  uint8_t s[32], p[32];

  init_base_result(&result, s, p);
  result.core_result.summary.status_code = 7;

  if (amduat_enc_pel1_result_encode_v1(&result, &encoded)) {
    fprintf(stderr, "expected encode failure for OK status with nonzero code\n");
    amduat_octets_free(&encoded);
    return 1;
  }

  return 0;
}

static int test_encode_store_failure_invariant(void) {
  amduat_pel_surface_execution_result_t result;
  amduat_octets_t encoded = amduat_octets(NULL, 0);
  uint8_t s[32], p[32], f[32];

  init_base_result(&result, s, p);
  result.has_store_failure = true;
  result.store_failure.phase = AMDUAT_PEL_STORE_FAILURE_PROGRAM;
  result.store_failure.error_code = AMDUAT_PEL_STORE_ERROR_NOT_FOUND;
  result.store_failure.failing_ref = make_ref(0x03, f);

  if (amduat_enc_pel1_result_encode_v1(&result, &encoded)) {
    fprintf(stderr, "expected encode failure for store_failure mismatch\n");
    amduat_octets_free(&encoded);
    return 1;
  }

  return 0;
}

static bool mutate_and_decode(amduat_octets_t encoded,
                              size_t offset,
                              uint8_t value) {
  amduat_pel_surface_execution_result_t decoded;
  bool ok;

  if (offset >= encoded.len) {
    return false;
  }
  ((uint8_t *)encoded.data)[offset] = value;
  ok = amduat_enc_pel1_result_decode_v1(encoded, &decoded);
  if (ok) {
    amduat_enc_pel1_result_free(&decoded);
  }
  return ok;
}

static int test_decode_invalid_status_kind(void) {
  amduat_pel_surface_execution_result_t result;
  amduat_octets_t encoded = amduat_octets(NULL, 0);
  uint8_t *original = NULL;
  uint8_t s[32], p[32];
  size_t encoded_ref_len = 4 + 2 + 32;
  size_t offset = 0;
  size_t status_offset;
  size_t kind_offset;

  init_base_result(&result, s, p);
  if (!amduat_enc_pel1_result_encode_v1(&result, &encoded)) {
    fprintf(stderr, "encode failed\n");
    return 1;
  }
  if (encoded.len != 0) {
    original = (uint8_t *)malloc(encoded.len);
    if (original == NULL) {
      fprintf(stderr, "alloc failed\n");
      amduat_octets_free(&encoded);
      return 1;
    }
    memcpy(original, encoded.data, encoded.len);
  }

  offset += 2;
  offset += encoded_ref_len;
  offset += encoded_ref_len;
  offset += 4;
  offset += 4;
  offset += 1;
  offset += 1;
  offset += 1;
  offset += 2;
  status_offset = offset;
  kind_offset = status_offset + 1 + encoded_ref_len;

  if (mutate_and_decode(encoded, status_offset, 0xffu)) {
    fprintf(stderr, "expected decode failure for invalid status\n");
    goto cleanup;
  }
  if (original != NULL) {
    memcpy((uint8_t *)encoded.data, original, encoded.len);
  }
  if (mutate_and_decode(encoded, kind_offset, 0xffu)) {
    fprintf(stderr, "expected decode failure for invalid summary kind\n");
    goto cleanup;
  }

cleanup:
  free(original);
  amduat_octets_free(&encoded);
  return 0;
}

int main(void) {
  if (test_encode_status_invariant() != 0) {
    return 1;
  }
  if (test_encode_store_failure_invariant() != 0) {
    return 1;
  }
  if (test_decode_invalid_status_kind() != 0) {
    return 1;
  }
  return 0;
}