#include "amduat/asl/ref_derive.h"
#include "amduat/asl/store.h"
#include "amduat/enc/asl1_core.h"
#include "amduat/enc/pel1_result.h"
#include "amduat/enc/pel_program_dag.h"
#include "amduat/enc/pel_trace_dag.h"
#include "amduat/hash/asl1.h"
#include "amduat/pel/opreg_kernel.h"
#include "amduat/pel/program_dag.h"
#include "amduat/pel/program_dag_desc.h"
#include "amduat/pel/surf.h"

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

typedef struct {
  amduat_reference_t ref;
  amduat_artifact_t artifact;
} stub_store_entry_t;

typedef struct {
  stub_store_entry_t *entries;
  size_t len;
  size_t cap;
  amduat_asl_store_config_t config;
  bool fail_get;
  amduat_reference_t fail_get_ref;
  amduat_asl_store_error_t fail_get_error;
} stub_store_t;

static void stub_store_init(stub_store_t *store) {
  if (store == NULL) {
    return;
  }
  store->entries = NULL;
  store->len = 0;
  store->cap = 0;
  store->config.encoding_profile_id = 0;
  store->config.hash_id = 0;
  store->fail_get = false;
  store->fail_get_ref = amduat_reference(0, amduat_octets(NULL, 0));
  store->fail_get_error = AMDUAT_ASL_STORE_OK;
}

static void stub_store_free(stub_store_t *store) {
  size_t i;

  if (store == NULL) {
    return;
  }
  for (i = 0; i < store->len; ++i) {
    stub_store_entry_t *entry = &store->entries[i];
    free((void *)entry->ref.digest.data);
    entry->ref.digest.data = NULL;
    entry->ref.digest.len = 0;
    free((void *)entry->artifact.bytes.data);
    entry->artifact.bytes.data = NULL;
    entry->artifact.bytes.len = 0;
  }
  free(store->entries);
  store->entries = NULL;
  store->len = 0;
  store->cap = 0;
  amduat_reference_free(&store->fail_get_ref);
  store->fail_get = false;
}

static bool artifact_eq(const amduat_artifact_t *a,
                        const amduat_artifact_t *b) {
  if (a == NULL || b == NULL) {
    return false;
  }
  if (a->has_type_tag != b->has_type_tag) {
    return false;
  }
  if (a->has_type_tag && a->type_tag.tag_id != b->type_tag.tag_id) {
    return false;
  }
  if (a->bytes.len != b->bytes.len) {
    return false;
  }
  if (a->bytes.len != 0) {
    if (a->bytes.data == NULL || b->bytes.data == NULL) {
      return false;
    }
    if (memcmp(a->bytes.data, b->bytes.data, a->bytes.len) != 0) {
      return false;
    }
  }
  return true;
}

static bool stub_store_add_entry(stub_store_t *store,
                                 amduat_reference_t stored_ref,
                                 amduat_artifact_t stored_artifact) {
  stub_store_entry_t *entries;
  size_t new_cap;

  if (store->len == store->cap) {
    new_cap = store->cap == 0 ? 8 : store->cap * 2;
    entries = (stub_store_entry_t *)realloc(store->entries,
                                            new_cap * sizeof(*entries));
    if (entries == NULL) {
      return false;
    }
    store->entries = entries;
    store->cap = new_cap;
  }

  store->entries[store->len++] =
      (stub_store_entry_t){stored_ref, stored_artifact};
  return true;
}

static amduat_asl_store_error_t stub_store_put(
    void *ctx,
    amduat_artifact_t artifact,
    amduat_reference_t *out_ref) {
  stub_store_t *store;
  const amduat_hash_asl1_desc_t *hash_desc;
  uint8_t *payload;
  amduat_reference_t derived_ref;
  amduat_reference_t stored_ref;
  amduat_artifact_t stored_artifact;
  size_t i;

  if (ctx == NULL || out_ref == NULL) {
    return AMDUAT_ASL_STORE_ERR_INTEGRITY;
  }
  store = (stub_store_t *)ctx;
  out_ref->hash_id = 0;
  out_ref->digest = amduat_octets(NULL, 0);

  if (store->config.encoding_profile_id != AMDUAT_ENC_ASL1_CORE_V1) {
    return AMDUAT_ASL_STORE_ERR_UNSUPPORTED;
  }
  hash_desc = amduat_hash_asl1_desc_lookup(store->config.hash_id);
  if (hash_desc == NULL || hash_desc->digest_len == 0) {
    return AMDUAT_ASL_STORE_ERR_UNSUPPORTED;
  }
  derived_ref = amduat_reference(0u, amduat_octets(NULL, 0u));
  if (!amduat_asl_ref_derive(artifact,
                             store->config.encoding_profile_id,
                             store->config.hash_id,
                             &derived_ref,
                             NULL)) {
    return AMDUAT_ASL_STORE_ERR_INTEGRITY;
  }
  *out_ref = derived_ref;

  for (i = 0; i < store->len; ++i) {
    stub_store_entry_t *entry = &store->entries[i];
    if (amduat_reference_eq(entry->ref, *out_ref)) {
      if (!artifact_eq(&entry->artifact, &artifact)) {
        amduat_reference_free(out_ref);
        return AMDUAT_ASL_STORE_ERR_INTEGRITY;
      }
      return AMDUAT_ASL_STORE_OK;
    }
  }

  if (!amduat_reference_clone(*out_ref, &stored_ref)) {
    amduat_reference_free(out_ref);
    return AMDUAT_ASL_STORE_ERR_INTEGRITY;
  }

  payload = NULL;
  if (artifact.bytes.len != 0) {
    if (artifact.bytes.data == NULL) {
      amduat_reference_free(out_ref);
      amduat_reference_free(&stored_ref);
      return AMDUAT_ASL_STORE_ERR_INTEGRITY;
    }
    payload = (uint8_t *)malloc(artifact.bytes.len);
    if (payload == NULL) {
      amduat_reference_free(out_ref);
      amduat_reference_free(&stored_ref);
      return AMDUAT_ASL_STORE_ERR_INTEGRITY;
    }
    memcpy(payload, artifact.bytes.data, artifact.bytes.len);
  }
  stored_artifact.bytes = amduat_octets(payload, artifact.bytes.len);
  stored_artifact.has_type_tag = artifact.has_type_tag;
  stored_artifact.type_tag = artifact.type_tag;

  if (!stub_store_add_entry(store, stored_ref, stored_artifact)) {
    free(payload);
    amduat_reference_free(out_ref);
    amduat_reference_free(&stored_ref);
    return AMDUAT_ASL_STORE_ERR_INTEGRITY;
  }

  return AMDUAT_ASL_STORE_OK;
}

static amduat_asl_store_error_t stub_store_get(
    void *ctx,
    amduat_reference_t ref,
    amduat_artifact_t *out_artifact) {
  stub_store_t *store;
  size_t i;

  if (ctx == NULL || out_artifact == NULL) {
    return AMDUAT_ASL_STORE_ERR_INTEGRITY;
  }
  store = (stub_store_t *)ctx;

  if (store->fail_get &&
      amduat_reference_eq(store->fail_get_ref, ref)) {
    return store->fail_get_error;
  }

  for (i = 0; i < store->len; ++i) {
    stub_store_entry_t *entry = &store->entries[i];
    if (amduat_reference_eq(entry->ref, ref)) {
      uint8_t *payload = NULL;
      if (entry->artifact.bytes.len != 0) {
        payload = (uint8_t *)malloc(entry->artifact.bytes.len);
        if (payload == NULL) {
          return AMDUAT_ASL_STORE_ERR_INTEGRITY;
        }
        memcpy(payload, entry->artifact.bytes.data, entry->artifact.bytes.len);
      }
      out_artifact->bytes = amduat_octets(payload,
                                          entry->artifact.bytes.len);
      out_artifact->has_type_tag = entry->artifact.has_type_tag;
      out_artifact->type_tag = entry->artifact.type_tag;
      return AMDUAT_ASL_STORE_OK;
    }
  }

  return AMDUAT_ASL_STORE_ERR_NOT_FOUND;
}

static void artifact_free(amduat_artifact_t *artifact) {
  if (artifact == NULL) {
    return;
  }
  free((void *)artifact->bytes.data);
  artifact->bytes.data = NULL;
  artifact->bytes.len = 0;
}

static void store_u64_be(uint8_t *out, uint64_t value) {
  out[0] = (uint8_t)((value >> 56) & 0xffu);
  out[1] = (uint8_t)((value >> 48) & 0xffu);
  out[2] = (uint8_t)((value >> 40) & 0xffu);
  out[3] = (uint8_t)((value >> 32) & 0xffu);
  out[4] = (uint8_t)((value >> 24) & 0xffu);
  out[5] = (uint8_t)((value >> 16) & 0xffu);
  out[6] = (uint8_t)((value >> 8) & 0xffu);
  out[7] = (uint8_t)(value & 0xffu);
}

static amduat_octets_t make_const_params(const uint8_t *bytes, size_t len) {
  size_t total = 1 + 8 + len;
  uint8_t *buffer = (uint8_t *)malloc(total);
  if (buffer == NULL) {
    return amduat_octets(NULL, 0);
  }
  buffer[0] = 0x00u;
  store_u64_be(buffer + 1, (uint64_t)len);
  if (len != 0 && bytes != NULL) {
    memcpy(buffer + 1 + 8, bytes, len);
  }
  return amduat_octets(buffer, total);
}

static bool build_const_program_artifact(amduat_artifact_t *out_artifact) {
  amduat_pel_node_t nodes[1];
  amduat_pel_root_ref_t roots[1];
  amduat_pel_program_t program;
  amduat_octets_t params;
  amduat_octets_t encoded;
  const char op_const[] = "pel.bytes.const";
  const uint8_t payload[] = {'h', 'i'};

  params = make_const_params(payload, sizeof(payload));
  if (params.data == NULL) {
    return false;
  }

  nodes[0].id = 1;
  nodes[0].op.name = amduat_octets(op_const, strlen(op_const));
  nodes[0].op.version = 1;
  nodes[0].inputs = NULL;
  nodes[0].inputs_len = 0;
  nodes[0].params = params;

  roots[0].node_id = 1;
  roots[0].output_index = 0;

  program.nodes = nodes;
  program.nodes_len = 1;
  program.roots = roots;
  program.roots_len = 1;

  encoded = amduat_octets(NULL, 0);
  if (!amduat_enc_pel_program_dag_encode_v1(&program, &encoded)) {
    free((void *)params.data);
    return false;
  }

  *out_artifact = amduat_artifact_with_type(
      encoded, amduat_type_tag(AMDUAT_PEL_TYPE_TAG_PROGRAM_DAG_1));
  free((void *)params.data);
  return true;
}

static bool build_concat_program_artifact(amduat_artifact_t *out_artifact) {
  amduat_pel_dag_input_t inputs[1];
  amduat_pel_node_t nodes[1];
  amduat_pel_root_ref_t roots[1];
  amduat_pel_program_t program;
  amduat_octets_t encoded;
  const char op_concat[] = "pel.bytes.concat";

  inputs[0].kind = AMDUAT_PEL_DAG_INPUT_EXTERNAL;
  inputs[0].value.external.input_index = 0;

  nodes[0].id = 1;
  nodes[0].op.name = amduat_octets(op_concat, strlen(op_concat));
  nodes[0].op.version = 1;
  nodes[0].inputs = inputs;
  nodes[0].inputs_len = 1;
  nodes[0].params = amduat_octets(NULL, 0);

  roots[0].node_id = 1;
  roots[0].output_index = 0;

  program.nodes = nodes;
  program.nodes_len = 1;
  program.roots = roots;
  program.roots_len = 1;

  encoded = amduat_octets(NULL, 0);
  if (!amduat_enc_pel_program_dag_encode_v1(&program, &encoded)) {
    return false;
  }

  *out_artifact = amduat_artifact_with_type(
      encoded, amduat_type_tag(AMDUAT_PEL_TYPE_TAG_PROGRAM_DAG_1));
  return true;
}

static bool build_concat2_program_artifact(amduat_artifact_t *out_artifact) {
  amduat_pel_dag_input_t inputs[2];
  amduat_pel_node_t nodes[1];
  amduat_pel_root_ref_t roots[1];
  amduat_pel_program_t program;
  amduat_octets_t encoded;
  const char op_concat[] = "pel.bytes.concat";

  inputs[0].kind = AMDUAT_PEL_DAG_INPUT_EXTERNAL;
  inputs[0].value.external.input_index = 0;
  inputs[1].kind = AMDUAT_PEL_DAG_INPUT_EXTERNAL;
  inputs[1].value.external.input_index = 1;

  nodes[0].id = 1;
  nodes[0].op.name = amduat_octets(op_concat, strlen(op_concat));
  nodes[0].op.version = 1;
  nodes[0].inputs = inputs;
  nodes[0].inputs_len = 2;
  nodes[0].params = amduat_octets(NULL, 0);

  roots[0].node_id = 1;
  roots[0].output_index = 0;

  program.nodes = nodes;
  program.nodes_len = 1;
  program.roots = roots;
  program.roots_len = 1;

  encoded = amduat_octets(NULL, 0);
  if (!amduat_enc_pel_program_dag_encode_v1(&program, &encoded)) {
    return false;
  }

  *out_artifact = amduat_artifact_with_type(
      encoded, amduat_type_tag(AMDUAT_PEL_TYPE_TAG_PROGRAM_DAG_1));
  return true;
}

static bool build_params_program_artifact(amduat_artifact_t *out_artifact) {
  amduat_pel_node_t nodes[1];
  amduat_pel_root_ref_t roots[1];
  amduat_pel_program_t program;
  amduat_octets_t encoded;
  const char op_params[] = "pel.bytes.params";

  nodes[0].id = 1;
  nodes[0].op.name = amduat_octets(op_params, strlen(op_params));
  nodes[0].op.version = 1;
  nodes[0].inputs = NULL;
  nodes[0].inputs_len = 0;
  nodes[0].params = amduat_octets(NULL, 0);

  roots[0].node_id = 1;
  roots[0].output_index = 0;

  program.nodes = nodes;
  program.nodes_len = 1;
  program.roots = roots;
  program.roots_len = 1;

  encoded = amduat_octets(NULL, 0);
  if (!amduat_enc_pel_program_dag_encode_v1(&program, &encoded)) {
    return false;
  }

  *out_artifact = amduat_artifact_with_type(
      encoded, amduat_type_tag(AMDUAT_PEL_TYPE_TAG_PROGRAM_DAG_1));
  return true;
}

static amduat_reference_t make_ref(uint8_t seed) {
  uint8_t *digest = (uint8_t *)malloc(32);
  if (digest == NULL) {
    return amduat_reference(0, amduat_octets(NULL, 0));
  }
  memset(digest, seed, 32);
  return amduat_reference(AMDUAT_HASH_ASL1_ID_SHA256,
                          amduat_octets(digest, 32));
}

static int test_surf_success(void) {
  stub_store_t stub;
  amduat_asl_store_t store;
  amduat_asl_store_ops_t ops;
  amduat_asl_store_config_t cfg;
  amduat_artifact_t program_artifact;
  amduat_reference_t program_ref;
  amduat_reference_t *output_refs = NULL;
  size_t output_refs_len = 0;
  amduat_reference_t result_ref;
  amduat_artifact_t result_artifact;
  amduat_pel_surface_execution_result_t decoded;
  amduat_artifact_t trace_artifact;
  amduat_pel_trace_dag_value_t trace;
  amduat_artifact_t exec_result_artifact;
  amduat_pel_surface_execution_result_t exec_result_decoded;
  const char op_const[] = "pel.bytes.const";
  int exit_code = 1;

  cfg.encoding_profile_id = AMDUAT_ENC_ASL1_CORE_V1;
  cfg.hash_id = AMDUAT_HASH_ASL1_ID_SHA256;
  stub_store_init(&stub);
  stub.config = cfg;
  amduat_asl_store_ops_init(&ops);
  ops.put = stub_store_put;
  ops.get = stub_store_get;
  amduat_asl_store_init(&store, cfg, ops, &stub);
  program_ref = amduat_reference(0, amduat_octets(NULL, 0));
  result_ref = amduat_reference(0, amduat_octets(NULL, 0));

  if (!build_const_program_artifact(&program_artifact)) {
    fprintf(stderr, "build program failed\n");
    goto cleanup_store;
  }
  if (stub_store_put(&stub, program_artifact, &program_ref) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "program put failed\n");
    free((void *)program_artifact.bytes.data);
    goto cleanup_store;
  }
  free((void *)program_artifact.bytes.data);

  if (!amduat_pel_surf_run(&store, amduat_pel_program_dag_scheme_ref(),
                           program_ref, NULL, 0, false,
                           amduat_reference(0, amduat_octets(NULL, 0)),
                           &output_refs, &output_refs_len, &result_ref)) {
    fprintf(stderr, "surf run failed\n");
    goto cleanup_refs;
  }

  if (output_refs_len != 1) {
    fprintf(stderr, "unexpected output ref count\n");
    goto cleanup_refs;
  }

  if (stub_store_get(&stub, result_ref, &result_artifact) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "result get failed\n");
    goto cleanup_refs;
  }
  if (!result_artifact.has_type_tag ||
      result_artifact.type_tag.tag_id != AMDUAT_TYPE_TAG_PEL1_RESULT_1) {
    fprintf(stderr, "result type tag mismatch\n");
    artifact_free(&result_artifact);
    goto cleanup_refs;
  }
  if (!amduat_enc_pel1_result_decode_v1(result_artifact.bytes, &decoded)) {
    fprintf(stderr, "result decode failed\n");
    artifact_free(&result_artifact);
    goto cleanup_refs;
  }
  artifact_free(&result_artifact);

  if (decoded.core_result.status != AMDUAT_PEL_EXEC_STATUS_OK ||
      decoded.core_result.summary.kind != AMDUAT_PEL_EXEC_ERROR_NONE ||
      decoded.core_result.summary.status_code != 0) {
    fprintf(stderr, "unexpected core status\n");
    goto cleanup_decoded;
  }
  if (!decoded.has_trace_ref) {
    fprintf(stderr, "missing trace ref\n");
    goto cleanup_decoded;
  }
  if (decoded.output_refs_len != 1 ||
      !amduat_reference_eq(decoded.output_refs[0], output_refs[0])) {
    fprintf(stderr, "output refs mismatch\n");
    goto cleanup_decoded;
  }

  if (stub_store_get(&stub, decoded.trace_ref, &trace_artifact) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "trace get failed\n");
    goto cleanup_decoded;
  }
  if (!trace_artifact.has_type_tag ||
      trace_artifact.type_tag.tag_id != AMDUAT_TYPE_TAG_PEL_TRACE_DAG_1) {
    fprintf(stderr, "trace type tag mismatch\n");
    artifact_free(&trace_artifact);
    goto cleanup_decoded;
  }
  if (!amduat_enc_pel_trace_dag_decode_v1(trace_artifact.bytes, &trace)) {
    fprintf(stderr, "trace decode failed\n");
    artifact_free(&trace_artifact);
    goto cleanup_decoded;
  }
  artifact_free(&trace_artifact);

  if (!trace.has_exec_result_ref) {
    fprintf(stderr, "missing trace exec_result_ref\n");
    amduat_enc_pel_trace_dag_free(&trace);
    goto cleanup_decoded;
  }
  if (stub_store_get(&stub, trace.exec_result_ref, &exec_result_artifact) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "exec result get failed\n");
    amduat_enc_pel_trace_dag_free(&trace);
    goto cleanup_decoded;
  }
  if (!exec_result_artifact.has_type_tag ||
      exec_result_artifact.type_tag.tag_id != AMDUAT_TYPE_TAG_PEL1_RESULT_1) {
    fprintf(stderr, "exec result type tag mismatch\n");
    artifact_free(&exec_result_artifact);
    amduat_enc_pel_trace_dag_free(&trace);
    goto cleanup_decoded;
  }
  if (!amduat_enc_pel1_result_decode_v1(exec_result_artifact.bytes,
                                        &exec_result_decoded)) {
    fprintf(stderr, "exec result decode failed\n");
    artifact_free(&exec_result_artifact);
    amduat_enc_pel_trace_dag_free(&trace);
    goto cleanup_decoded;
  }
  artifact_free(&exec_result_artifact);
  amduat_enc_pel1_result_free(&exec_result_decoded);

  if (trace.node_traces_len != 1 ||
      trace.node_traces[0].status != AMDUAT_PEL_NODE_TRACE_OK ||
      trace.node_traces[0].output_refs_len != 1 ||
      !amduat_reference_eq(trace.node_traces[0].output_refs[0],
                           output_refs[0]) ||
      trace.node_traces[0].op_name.len != strlen(op_const) ||
      memcmp(trace.node_traces[0].op_name.data,
             op_const, trace.node_traces[0].op_name.len) != 0) {
    fprintf(stderr, "trace node mismatch\n");
    amduat_enc_pel_trace_dag_free(&trace);
    goto cleanup_decoded;
  }

  amduat_enc_pel_trace_dag_free(&trace);
  exit_code = 0;

cleanup_decoded:
  amduat_enc_pel1_result_free(&decoded);
cleanup_refs:
  amduat_pel_surf_free_refs(output_refs, output_refs_len);
  amduat_pel_surf_free_ref(&result_ref);
  amduat_pel_surf_free_ref(&program_ref);
cleanup_store:
  stub_store_free(&stub);
  return exit_code;
}

static int test_surf_missing_program(void) {
  stub_store_t stub;
  amduat_asl_store_t store;
  amduat_asl_store_ops_t ops;
  amduat_asl_store_config_t cfg;
  amduat_reference_t program_ref;
  amduat_reference_t *output_refs = NULL;
  size_t output_refs_len = 0;
  amduat_reference_t result_ref;
  amduat_artifact_t result_artifact;
  amduat_pel_surface_execution_result_t decoded;
  int exit_code = 1;

  cfg.encoding_profile_id = AMDUAT_ENC_ASL1_CORE_V1;
  cfg.hash_id = AMDUAT_HASH_ASL1_ID_SHA256;
  stub_store_init(&stub);
  stub.config = cfg;
  amduat_asl_store_ops_init(&ops);
  ops.put = stub_store_put;
  ops.get = stub_store_get;
  amduat_asl_store_init(&store, cfg, ops, &stub);

  program_ref = make_ref(0x42u);
  if (program_ref.digest.data == NULL) {
    goto cleanup_store;
  }

  if (!amduat_pel_surf_run(&store, amduat_pel_program_dag_scheme_ref(),
                           program_ref, NULL, 0, false,
                           amduat_reference(0, amduat_octets(NULL, 0)),
                           &output_refs, &output_refs_len, &result_ref)) {
    fprintf(stderr, "surf run failed\n");
    goto cleanup_refs;
  }

  if (output_refs_len != 0) {
    fprintf(stderr, "unexpected output refs\n");
    goto cleanup_refs;
  }

  if (stub_store_get(&stub, result_ref, &result_artifact) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "result get failed\n");
    goto cleanup_refs;
  }
  if (!amduat_enc_pel1_result_decode_v1(result_artifact.bytes, &decoded)) {
    artifact_free(&result_artifact);
    fprintf(stderr, "result decode failed\n");
    goto cleanup_refs;
  }
  artifact_free(&result_artifact);

  if (decoded.core_result.status != AMDUAT_PEL_EXEC_STATUS_INVALID_PROGRAM ||
      decoded.core_result.summary.kind != AMDUAT_PEL_EXEC_ERROR_PROGRAM ||
      decoded.core_result.summary.status_code != 2 ||
      !decoded.has_store_failure ||
      decoded.store_failure.phase != AMDUAT_PEL_STORE_FAILURE_PROGRAM ||
      decoded.store_failure.error_code != AMDUAT_PEL_STORE_ERROR_NOT_FOUND ||
      !amduat_reference_eq(decoded.store_failure.failing_ref, program_ref) ||
      decoded.has_trace_ref) {
    fprintf(stderr, "missing program mapping mismatch\n");
    goto cleanup_decoded;
  }

  exit_code = 0;

cleanup_decoded:
  amduat_enc_pel1_result_free(&decoded);
cleanup_refs:
  amduat_pel_surf_free_refs(output_refs, output_refs_len);
  amduat_pel_surf_free_ref(&result_ref);
  amduat_pel_surf_free_ref(&program_ref);
cleanup_store:
  stub_store_free(&stub);
  return exit_code;
}

static int test_surf_missing_input(void) {
  stub_store_t stub;
  amduat_asl_store_t store;
  amduat_asl_store_ops_t ops;
  amduat_asl_store_config_t cfg;
  amduat_artifact_t program_artifact;
  amduat_reference_t program_ref;
  amduat_reference_t input_ref;
  amduat_reference_t *output_refs = NULL;
  size_t output_refs_len = 0;
  amduat_reference_t result_ref;
  amduat_artifact_t result_artifact;
  amduat_pel_surface_execution_result_t decoded;
  int exit_code = 1;

  cfg.encoding_profile_id = AMDUAT_ENC_ASL1_CORE_V1;
  cfg.hash_id = AMDUAT_HASH_ASL1_ID_SHA256;
  stub_store_init(&stub);
  stub.config = cfg;
  amduat_asl_store_ops_init(&ops);
  ops.put = stub_store_put;
  ops.get = stub_store_get;
  amduat_asl_store_init(&store, cfg, ops, &stub);

  if (!build_concat_program_artifact(&program_artifact)) {
    goto cleanup_store;
  }
  if (stub_store_put(&stub, program_artifact, &program_ref) !=
      AMDUAT_ASL_STORE_OK) {
    free((void *)program_artifact.bytes.data);
    goto cleanup_store;
  }
  free((void *)program_artifact.bytes.data);

  input_ref = make_ref(0x33u);
  if (input_ref.digest.data == NULL) {
    goto cleanup_refs;
  }

  if (!amduat_pel_surf_run(&store, amduat_pel_program_dag_scheme_ref(),
                           program_ref, &input_ref, 1, false,
                           amduat_reference(0, amduat_octets(NULL, 0)),
                           &output_refs, &output_refs_len, &result_ref)) {
    fprintf(stderr, "surf run failed\n");
    goto cleanup_refs;
  }

  if (output_refs_len != 0) {
    fprintf(stderr, "unexpected output refs\n");
    goto cleanup_refs;
  }

  if (stub_store_get(&stub, result_ref, &result_artifact) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "result get failed\n");
    goto cleanup_refs;
  }
  if (!amduat_enc_pel1_result_decode_v1(result_artifact.bytes, &decoded)) {
    artifact_free(&result_artifact);
    fprintf(stderr, "result decode failed\n");
    goto cleanup_refs;
  }
  artifact_free(&result_artifact);

  if (decoded.core_result.status != AMDUAT_PEL_EXEC_STATUS_INVALID_INPUTS ||
      decoded.core_result.summary.kind != AMDUAT_PEL_EXEC_ERROR_INPUTS ||
      decoded.core_result.summary.status_code != 3 ||
      !decoded.has_store_failure ||
      decoded.store_failure.phase != AMDUAT_PEL_STORE_FAILURE_INPUT ||
      decoded.store_failure.error_code != AMDUAT_PEL_STORE_ERROR_NOT_FOUND ||
      !amduat_reference_eq(decoded.store_failure.failing_ref, input_ref) ||
      decoded.has_trace_ref) {
    fprintf(stderr, "missing input mapping mismatch\n");
    goto cleanup_decoded;
  }

  exit_code = 0;

cleanup_decoded:
  amduat_enc_pel1_result_free(&decoded);
cleanup_refs:
  amduat_pel_surf_free_refs(output_refs, output_refs_len);
  amduat_pel_surf_free_ref(&result_ref);
  amduat_pel_surf_free_ref(&program_ref);
  amduat_pel_surf_free_ref(&input_ref);
cleanup_store:
  stub_store_free(&stub);
  return exit_code;
}

static int test_surf_runtime_failure_trace_diag(void) {
  stub_store_t stub;
  amduat_asl_store_t store;
  amduat_asl_store_ops_t ops;
  amduat_asl_store_config_t cfg;
  amduat_artifact_t program_artifact;
  amduat_reference_t program_ref;
  amduat_reference_t input_refs[2];
  amduat_reference_t *output_refs = NULL;
  size_t output_refs_len = 0;
  amduat_reference_t result_ref;
  amduat_artifact_t result_artifact;
  amduat_pel_surface_execution_result_t decoded;
  amduat_artifact_t trace_artifact;
  amduat_pel_trace_dag_value_t trace;
  int exit_code = 1;
  uint8_t payload_a[] = {'a'};
  uint8_t payload_b[] = {'b'};
  amduat_artifact_t input_a;
  amduat_artifact_t input_b;

  cfg.encoding_profile_id = AMDUAT_ENC_ASL1_CORE_V1;
  cfg.hash_id = AMDUAT_HASH_ASL1_ID_SHA256;
  stub_store_init(&stub);
  stub.config = cfg;
  amduat_asl_store_ops_init(&ops);
  ops.put = stub_store_put;
  ops.get = stub_store_get;
  amduat_asl_store_init(&store, cfg, ops, &stub);

  if (!build_concat2_program_artifact(&program_artifact)) {
    fprintf(stderr, "build program failed\n");
    goto cleanup_store;
  }
  if (stub_store_put(&stub, program_artifact, &program_ref) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "program put failed\n");
    free((void *)program_artifact.bytes.data);
    goto cleanup_store;
  }
  free((void *)program_artifact.bytes.data);

  input_a = amduat_artifact(amduat_octets(payload_a, sizeof(payload_a)));
  input_b = amduat_artifact_with_type(
      amduat_octets(payload_b, sizeof(payload_b)),
      amduat_type_tag(AMDUAT_TYPE_TAG_PEL1_RESULT_1));

  if (stub_store_put(&stub, input_a, &input_refs[0]) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "input a put failed\n");
    goto cleanup_store;
  }
  if (stub_store_put(&stub, input_b, &input_refs[1]) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "input b put failed\n");
    goto cleanup_store;
  }

  if (!amduat_pel_surf_run(&store, amduat_pel_program_dag_scheme_ref(),
                           program_ref, input_refs, 2, false,
                           amduat_reference(0, amduat_octets(NULL, 0)),
                           &output_refs, &output_refs_len, &result_ref)) {
    fprintf(stderr, "surf run failed\n");
    goto cleanup_refs;
  }

  if (output_refs_len != 0) {
    fprintf(stderr, "unexpected output refs\n");
    goto cleanup_refs;
  }

  if (stub_store_get(&stub, result_ref, &result_artifact) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "result get failed\n");
    goto cleanup_refs;
  }
  if (!amduat_enc_pel1_result_decode_v1(result_artifact.bytes, &decoded)) {
    artifact_free(&result_artifact);
    fprintf(stderr, "result decode failed\n");
    goto cleanup_refs;
  }
  artifact_free(&result_artifact);

  if (decoded.core_result.status != AMDUAT_PEL_EXEC_STATUS_RUNTIME_FAILED ||
      decoded.core_result.summary.kind != AMDUAT_PEL_EXEC_ERROR_RUNTIME ||
      decoded.core_result.summary.status_code == 0 ||
      !decoded.has_trace_ref) {
    fprintf(stderr, "runtime result mismatch\n");
    goto cleanup_decoded;
  }

  if (stub_store_get(&stub, decoded.trace_ref, &trace_artifact) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "trace get failed\n");
    goto cleanup_decoded;
  }
  if (!amduat_enc_pel_trace_dag_decode_v1(trace_artifact.bytes, &trace)) {
    fprintf(stderr, "trace decode failed\n");
    artifact_free(&trace_artifact);
    goto cleanup_decoded;
  }
  artifact_free(&trace_artifact);

  if (trace.node_traces_len != 1 ||
      trace.node_traces[0].status != AMDUAT_PEL_NODE_TRACE_FAILED ||
      trace.node_traces[0].status_code !=
          AMDUAT_PEL_KERNEL_STATUS_CONCAT_TYPE_TAG_MISMATCH ||
      trace.node_traces[0].diagnostics_len != 0) {
    fprintf(stderr, "trace runtime diagnostics mismatch\n");
    amduat_enc_pel_trace_dag_free(&trace);
    goto cleanup_decoded;
  }

  amduat_enc_pel_trace_dag_free(&trace);
  exit_code = 0;

cleanup_decoded:
  amduat_enc_pel1_result_free(&decoded);
cleanup_refs:
  amduat_pel_surf_free_refs(output_refs, output_refs_len);
  amduat_pel_surf_free_ref(&result_ref);
  amduat_pel_surf_free_ref(&program_ref);
cleanup_store:
  stub_store_free(&stub);
  return exit_code;
}

static int test_surf_store_io_failure(void) {
  stub_store_t stub;
  amduat_asl_store_t store;
  amduat_asl_store_ops_t ops;
  amduat_asl_store_config_t cfg;
  amduat_artifact_t program_artifact;
  amduat_reference_t program_ref;
  amduat_reference_t *output_refs = NULL;
  size_t output_refs_len = 0;
  amduat_reference_t result_ref;
  size_t entries_before = 0;
  int exit_code = 1;

  cfg.encoding_profile_id = AMDUAT_ENC_ASL1_CORE_V1;
  cfg.hash_id = AMDUAT_HASH_ASL1_ID_SHA256;
  stub_store_init(&stub);
  stub.config = cfg;
  amduat_asl_store_ops_init(&ops);
  ops.put = stub_store_put;
  ops.get = stub_store_get;
  amduat_asl_store_init(&store, cfg, ops, &stub);

  if (!build_const_program_artifact(&program_artifact)) {
    fprintf(stderr, "build program failed\n");
    goto cleanup_store;
  }
  if (stub_store_put(&stub, program_artifact, &program_ref) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "program put failed\n");
    free((void *)program_artifact.bytes.data);
    goto cleanup_store;
  }
  free((void *)program_artifact.bytes.data);

  if (!amduat_reference_clone(program_ref, &stub.fail_get_ref)) {
    fprintf(stderr, "fail ref clone failed\n");
    goto cleanup_refs;
  }
  stub.fail_get = true;
  stub.fail_get_error = AMDUAT_ASL_STORE_ERR_IO;
  entries_before = stub.len;

  if (amduat_pel_surf_run(&store, amduat_pel_program_dag_scheme_ref(),
                          program_ref, NULL, 0, false,
                          amduat_reference(0, amduat_octets(NULL, 0)),
                          &output_refs, &output_refs_len, &result_ref)) {
    fprintf(stderr, "expected surf run failure\n");
    goto cleanup_refs;
  }

  if (stub.len != entries_before) {
    fprintf(stderr, "unexpected store mutation\n");
    goto cleanup_refs;
  }

  exit_code = 0;

cleanup_refs:
  amduat_pel_surf_free_refs(output_refs, output_refs_len);
  amduat_pel_surf_free_ref(&result_ref);
  amduat_pel_surf_free_ref(&program_ref);
cleanup_store:
  stub_store_free(&stub);
  return exit_code;
}

static int test_surf_params_pass_through(void) {
  stub_store_t stub;
  amduat_asl_store_t store;
  amduat_asl_store_ops_t ops;
  amduat_asl_store_config_t cfg;
  amduat_artifact_t program_artifact;
  amduat_reference_t program_ref;
  amduat_artifact_t params_artifact;
  amduat_reference_t params_ref;
  amduat_reference_t *output_refs = NULL;
  size_t output_refs_len = 0;
  amduat_reference_t result_ref;
  amduat_artifact_t output_artifact;
  amduat_artifact_t result_artifact;
  amduat_pel_surface_execution_result_t decoded;
  const uint8_t params_bytes[] = {'p', 'a', 'r', 'a', 'm', 's'};
  int exit_code = 1;

  cfg.encoding_profile_id = AMDUAT_ENC_ASL1_CORE_V1;
  cfg.hash_id = AMDUAT_HASH_ASL1_ID_SHA256;
  stub_store_init(&stub);
  stub.config = cfg;
  amduat_asl_store_ops_init(&ops);
  ops.put = stub_store_put;
  ops.get = stub_store_get;
  amduat_asl_store_init(&store, cfg, ops, &stub);
  program_ref = amduat_reference(0, amduat_octets(NULL, 0));
  params_ref = amduat_reference(0, amduat_octets(NULL, 0));
  result_ref = amduat_reference(0, amduat_octets(NULL, 0));

  if (!build_params_program_artifact(&program_artifact)) {
    fprintf(stderr, "build params program failed\n");
    goto cleanup_store;
  }
  if (stub_store_put(&stub, program_artifact, &program_ref) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "program put failed\n");
    free((void *)program_artifact.bytes.data);
    goto cleanup_store;
  }
  free((void *)program_artifact.bytes.data);

  params_artifact = amduat_artifact(amduat_octets(params_bytes,
                                                  sizeof(params_bytes)));
  if (stub_store_put(&stub, params_artifact, &params_ref) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "params put failed\n");
    goto cleanup_refs;
  }

  if (!amduat_pel_surf_run(&store, amduat_pel_program_dag_scheme_ref(),
                           program_ref, NULL, 0, true, params_ref,
                           &output_refs, &output_refs_len, &result_ref)) {
    fprintf(stderr, "surf run failed\n");
    goto cleanup_refs;
  }

  if (output_refs_len != 1) {
    fprintf(stderr, "unexpected output ref count\n");
    goto cleanup_refs;
  }

  if (stub_store_get(&stub, output_refs[0], &output_artifact) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "output get failed\n");
    goto cleanup_refs;
  }
  if (!artifact_eq(&output_artifact, &params_artifact)) {
    artifact_free(&output_artifact);
    fprintf(stderr, "params output mismatch\n");
    goto cleanup_refs;
  }
  artifact_free(&output_artifact);

  if (stub_store_get(&stub, result_ref, &result_artifact) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "result get failed\n");
    goto cleanup_refs;
  }
  if (!amduat_enc_pel1_result_decode_v1(result_artifact.bytes, &decoded)) {
    artifact_free(&result_artifact);
    fprintf(stderr, "result decode failed\n");
    goto cleanup_refs;
  }
  artifact_free(&result_artifact);

  if (!decoded.has_params_ref ||
      !amduat_reference_eq(decoded.params_ref, params_ref)) {
    fprintf(stderr, "missing params ref\n");
    goto cleanup_decoded;
  }

  exit_code = 0;

cleanup_decoded:
  amduat_enc_pel1_result_free(&decoded);
cleanup_refs:
  amduat_pel_surf_free_refs(output_refs, output_refs_len);
  amduat_pel_surf_free_ref(&result_ref);
  amduat_pel_surf_free_ref(&program_ref);
  amduat_pel_surf_free_ref(&params_ref);
cleanup_store:
  stub_store_free(&stub);
  return exit_code;
}

static int test_surf_missing_params(void) {
  stub_store_t stub;
  amduat_asl_store_t store;
  amduat_asl_store_ops_t ops;
  amduat_asl_store_config_t cfg;
  amduat_artifact_t program_artifact;
  amduat_reference_t program_ref;
  amduat_reference_t *output_refs = NULL;
  size_t output_refs_len = 0;
  amduat_reference_t result_ref;
  amduat_artifact_t result_artifact;
  amduat_pel_surface_execution_result_t decoded;
  int exit_code = 1;

  cfg.encoding_profile_id = AMDUAT_ENC_ASL1_CORE_V1;
  cfg.hash_id = AMDUAT_HASH_ASL1_ID_SHA256;
  stub_store_init(&stub);
  stub.config = cfg;
  amduat_asl_store_ops_init(&ops);
  ops.put = stub_store_put;
  ops.get = stub_store_get;
  amduat_asl_store_init(&store, cfg, ops, &stub);
  program_ref = amduat_reference(0, amduat_octets(NULL, 0));
  result_ref = amduat_reference(0, amduat_octets(NULL, 0));

  if (!build_params_program_artifact(&program_artifact)) {
    fprintf(stderr, "build params program failed\n");
    goto cleanup_store;
  }
  if (stub_store_put(&stub, program_artifact, &program_ref) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "program put failed\n");
    free((void *)program_artifact.bytes.data);
    goto cleanup_store;
  }
  free((void *)program_artifact.bytes.data);

  if (!amduat_pel_surf_run(&store, amduat_pel_program_dag_scheme_ref(),
                           program_ref, NULL, 0, false,
                           amduat_reference(0, amduat_octets(NULL, 0)),
                           &output_refs, &output_refs_len, &result_ref)) {
    fprintf(stderr, "surf run failed\n");
    goto cleanup_refs;
  }

  if (output_refs_len != 0) {
    fprintf(stderr, "unexpected output refs\n");
    goto cleanup_refs;
  }

  if (stub_store_get(&stub, result_ref, &result_artifact) !=
      AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "result get failed\n");
    goto cleanup_refs;
  }
  if (!amduat_enc_pel1_result_decode_v1(result_artifact.bytes, &decoded)) {
    artifact_free(&result_artifact);
    fprintf(stderr, "result decode failed\n");
    goto cleanup_refs;
  }
  artifact_free(&result_artifact);

  if (decoded.core_result.status != AMDUAT_PEL_EXEC_STATUS_INVALID_INPUTS ||
      decoded.core_result.summary.kind != AMDUAT_PEL_EXEC_ERROR_INPUTS ||
      decoded.core_result.summary.status_code != 3 ||
      decoded.has_store_failure ||
      decoded.has_params_ref) {
    fprintf(stderr, "missing params mapping mismatch\n");
    goto cleanup_decoded;
  }

  exit_code = 0;

cleanup_decoded:
  amduat_enc_pel1_result_free(&decoded);
cleanup_refs:
  amduat_pel_surf_free_refs(output_refs, output_refs_len);
  amduat_pel_surf_free_ref(&result_ref);
  amduat_pel_surf_free_ref(&program_ref);
cleanup_store:
  stub_store_free(&stub);
  return exit_code;
}

int main(void) {
  if (test_surf_success() != 0) {
    return 1;
  }
  if (test_surf_missing_program() != 0) {
    return 1;
  }
  if (test_surf_missing_input() != 0) {
    return 1;
  }
  if (test_surf_runtime_failure_trace_diag() != 0) {
    return 1;
  }
  if (test_surf_store_io_failure() != 0) {
    return 1;
  }
  if (test_surf_params_pass_through() != 0) {
    return 1;
  }
  if (test_surf_missing_params() != 0) {
    return 1;
  }
  return 0;
}