#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/hash/asl1.h"
#include "amduat/pel/program_dag.h"
#include "amduat/pel/program_dag_desc.h"
#include "amduat/pel/run.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;
} 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 = AMDUAT_ENC_ASL1_CORE_V1;
  store->config.hash_id = AMDUAT_HASH_ASL1_ID_SHA256;
}

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;
}

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;
  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);

  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) {
    if (amduat_reference_eq(store->entries[i].ref, *out_ref)) {
      return AMDUAT_ASL_STORE_OK;
    }
  }

  if (!amduat_reference_clone(*out_ref, &stored_ref)) {
    amduat_reference_free(out_ref);
    return AMDUAT_ASL_STORE_ERR_INTEGRITY;
  }
  stored_artifact.bytes = amduat_octets(NULL, 0u);
  if (artifact.bytes.len != 0) {
    uint8_t *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)) {
    amduat_reference_free(out_ref);
    amduat_reference_free(&stored_ref);
    free((void *)stored_artifact.bytes.data);
    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;
  out_artifact->bytes = amduat_octets(NULL, 0u);
  out_artifact->has_type_tag = false;
  out_artifact->type_tag = amduat_type_tag(0);

  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 != 0u) {
        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 int test_program_concat_two_inputs(void) {
  stub_store_t stub;
  amduat_asl_store_ops_t ops;
  amduat_asl_store_t store;
  amduat_pel_program_t program;
  amduat_pel_node_t *nodes;
  amduat_pel_dag_input_t *inputs;
  amduat_pel_root_ref_t *roots;
  amduat_octets_t encoded;
  amduat_type_tag_t program_tag;
  amduat_asl_encoding_profile_id_t program_profile;
  amduat_artifact_t program_artifact;
  amduat_reference_t program_ref;
  amduat_artifact_t input_artifacts[2];
  amduat_reference_t input_refs[2];
  amduat_pel_run_result_t run_result;
  amduat_artifact_t output_artifact;
  int exit_code = 0;

  stub_store_init(&stub);
  amduat_asl_store_ops_init(&ops);
  ops.put = stub_store_put;
  ops.get = stub_store_get;
  amduat_asl_store_init(&store, stub.config, ops, &stub);

  nodes = (amduat_pel_node_t *)calloc(1u, sizeof(*nodes));
  inputs = (amduat_pel_dag_input_t *)calloc(2u, sizeof(*inputs));
  roots = (amduat_pel_root_ref_t *)calloc(1u, sizeof(*roots));
  if (nodes == NULL || inputs == NULL || roots == NULL) {
    fprintf(stderr, "out of memory\n");
    exit_code = 1;
    goto cleanup;
  }

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

  nodes[0].id = 1u;
  nodes[0].op.name =
      amduat_octets("pel.bytes.concat", strlen("pel.bytes.concat"));
  nodes[0].op.version = 1u;
  nodes[0].inputs = inputs;
  nodes[0].inputs_len = 2u;
  nodes[0].params = amduat_octets(NULL, 0u);

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

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

  if (!amduat_pel_program_dag_validate(&program)) {
    fprintf(stderr, "program validation failed\n");
    exit_code = 1;
    goto cleanup;
  }

  encoded = amduat_octets(NULL, 0u);
  if (!amduat_enc_pel_program_dag_encode_v1(&program, &encoded)) {
    fprintf(stderr, "program encode failed\n");
    exit_code = 1;
    goto cleanup;
  }

  if (!amduat_pel_program_dag_desc_get_program_binding(&program_tag,
                                                       &program_profile)) {
    fprintf(stderr, "program binding failed\n");
    free((void *)encoded.data);
    exit_code = 1;
    goto cleanup;
  }
  (void)program_profile;

  program_artifact = amduat_artifact_with_type(encoded, program_tag);
  program_ref = amduat_reference(0u, amduat_octets(NULL, 0u));
  if (amduat_asl_store_put(&store, program_artifact,
                           &program_ref) != AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "store put program failed\n");
    free((void *)encoded.data);
    exit_code = 1;
    goto cleanup;
  }
  free((void *)encoded.data);

  input_artifacts[0] =
      amduat_artifact(amduat_octets("foo", 3u));
  input_artifacts[1] =
      amduat_artifact(amduat_octets("bar", 3u));

  input_refs[0] = amduat_reference(0u, amduat_octets(NULL, 0u));
  input_refs[1] = amduat_reference(0u, amduat_octets(NULL, 0u));
  if (amduat_asl_store_put(&store, input_artifacts[0],
                           &input_refs[0]) != AMDUAT_ASL_STORE_OK ||
      amduat_asl_store_put(&store, input_artifacts[1],
                           &input_refs[1]) != AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "store put inputs failed\n");
    exit_code = 1;
    goto cleanup;
  }

  memset(&run_result, 0, sizeof(run_result));
  if (!amduat_pel_surf_run_with_result(
          &store,
          amduat_pel_program_dag_scheme_ref(),
          program_ref,
          input_refs,
          2u,
          false,
          amduat_reference(0u, amduat_octets(NULL, 0u)),
          &run_result)) {
    fprintf(stderr, "surf run failed\n");
    exit_code = 1;
    goto cleanup;
  }

  if (!run_result.has_result_value ||
      run_result.result_value.core_result.status != AMDUAT_PEL_EXEC_STATUS_OK) {
    fprintf(stderr, "execution status not OK\n");
    exit_code = 1;
    goto cleanup_run;
  }
  if (run_result.output_refs_len != 1u) {
    fprintf(stderr, "unexpected output_refs_len=%zu\n",
            run_result.output_refs_len);
    exit_code = 1;
    goto cleanup_run;
  }

  memset(&output_artifact, 0, sizeof(output_artifact));
  if (amduat_asl_store_get(&store,
                           run_result.output_refs[0],
                           &output_artifact) != AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "store get output failed\n");
    exit_code = 1;
    goto cleanup_run;
  }
  if (output_artifact.bytes.len != 6u ||
      memcmp(output_artifact.bytes.data, "foobar", 6u) != 0) {
    fprintf(stderr, "unexpected output bytes\n");
    exit_code = 1;
  }
  free((void *)output_artifact.bytes.data);

cleanup_run:
  if (run_result.has_result_value) {
    amduat_enc_pel1_result_free(&run_result.result_value);
  }
  if (run_result.output_refs != NULL) {
    amduat_pel_surf_free_refs(run_result.output_refs,
                              run_result.output_refs_len);
  }
  amduat_pel_surf_free_ref(&run_result.result_ref);

cleanup:
  free(nodes);
  free(inputs);
  free(roots);
  stub_store_free(&stub);
  return exit_code;
}

int main(void) {
  return test_program_concat_two_inputs();
}