#include "amduat/enc/asl1_core.h"
#include "amduat/enc/tgk1_edge.h"
#include "amduat/hash/asl1.h"
#include "amduat/tgk/tgk_store_mem.h"

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

static const uint8_t k_non_edge_bytes[] = {0xde, 0xad, 0xbe, 0xef};

typedef struct {
  amduat_tgk_store_t store;
  amduat_tgk_store_mem_t mem;
  amduat_tgk_store_mem_artifact_t artifacts[5];
  amduat_octets_t edge_bytes[4];
  amduat_reference_t ref_edge1;
  amduat_reference_t ref_edge2;
  amduat_reference_t ref_edge3;
  amduat_reference_t ref_edge_bad;
  amduat_reference_t ref_non_edge;
  amduat_reference_t ref_missing;
  amduat_reference_t node_a;
  amduat_reference_t node_b;
  amduat_reference_t node_c;
  amduat_reference_t payload1;
  amduat_reference_t payload2;
  amduat_reference_t payload3;
  amduat_reference_t payload_bad;
  uint8_t digest_edge1[32];
  uint8_t digest_edge2[32];
  uint8_t digest_edge3[32];
  uint8_t digest_edge_bad[32];
  uint8_t digest_non_edge[32];
  uint8_t digest_missing[32];
  uint8_t digest_node_a[32];
  uint8_t digest_node_b[32];
  uint8_t digest_node_c[32];
  uint8_t digest_payload1[32];
  uint8_t digest_payload2[32];
  uint8_t digest_payload3[32];
  uint8_t digest_payload_bad[32];
  amduat_tgk_identity_domain_t domains[1];
  uint32_t edge_tags[1];
  amduat_tgk_edge_type_id_t edge_types[2];
  amduat_asl_encoding_profile_id_t encodings[1];
  amduat_tgk_store_config_t config;
} test_env_t;

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(AMDUAT_HASH_ASL1_ID_SHA256,
                          amduat_octets(storage, 32));
}

static void free_edge_bytes(test_env_t *env) {
  size_t i;

  for (i = 0; i < 4; ++i) {
    free((void *)env->edge_bytes[i].data);
    env->edge_bytes[i].data = NULL;
    env->edge_bytes[i].len = 0;
  }
}

static bool init_env(test_env_t *env) {
  amduat_tgk_edge_body_t edge1;
  amduat_tgk_edge_body_t edge2;
  amduat_tgk_edge_body_t edge3;
  amduat_tgk_edge_body_t edge_bad;
  amduat_reference_t edge1_from[1];
  amduat_reference_t edge1_to[1];
  amduat_reference_t edge2_from[1];
  amduat_reference_t edge2_to[1];
  amduat_reference_t edge3_from[1];
  amduat_reference_t edge3_to[1];
  amduat_reference_t edge_bad_from[1];
  amduat_reference_t edge_bad_to[1];

  memset(env, 0, sizeof(*env));

  env->domains[0].encoding_profile = AMDUAT_ENC_ASL1_CORE_V1;
  env->domains[0].hash_id = AMDUAT_HASH_ASL1_ID_SHA256;
  env->edge_tags[0] = TYPE_TAG_TGK1_EDGE_V1;
  env->edge_types[0] = 0x10;
  env->edge_types[1] = 0x20;
  env->encodings[0] = TGK1_EDGE_ENC_V1;

  env->config.id_space.domains = env->domains;
  env->config.id_space.domains_len = 1;
  env->config.artifact_scope.description = amduat_octets(NULL, 0);
  env->config.tgk_profiles.edge_tags = env->edge_tags;
  env->config.tgk_profiles.edge_tags_len = 1;
  env->config.tgk_profiles.edge_types = env->edge_types;
  env->config.tgk_profiles.edge_types_len = 2;
  env->config.tgk_profiles.encodings = env->encodings;
  env->config.tgk_profiles.encodings_len = 1;

  env->ref_edge1 = make_ref(0x20, env->digest_edge1);
  env->ref_edge2 = make_ref(0x10, env->digest_edge2);
  env->ref_edge3 = make_ref(0x30, env->digest_edge3);
  env->ref_edge_bad = make_ref(0x40, env->digest_edge_bad);
  env->ref_non_edge = make_ref(0x50, env->digest_non_edge);
  env->ref_missing = make_ref(0x60, env->digest_missing);

  env->node_a = make_ref(0xa1, env->digest_node_a);
  env->node_b = make_ref(0xb1, env->digest_node_b);
  env->node_c = make_ref(0xc1, env->digest_node_c);

  env->payload1 = make_ref(0xe1, env->digest_payload1);
  env->payload2 = make_ref(0xe2, env->digest_payload2);
  env->payload3 = make_ref(0xe3, env->digest_payload3);
  env->payload_bad = make_ref(0xe4, env->digest_payload_bad);

  memset(&edge1, 0, sizeof(edge1));
  edge1.type = 0x10;
  edge1_from[0] = env->node_a;
  edge1.from = edge1_from;
  edge1.from_len = 1;
  edge1_to[0] = env->node_b;
  edge1.to = edge1_to;
  edge1.to_len = 1;
  edge1.payload = env->payload1;

  memset(&edge2, 0, sizeof(edge2));
  edge2.type = 0x10;
  edge2_from[0] = env->node_a;
  edge2.from = edge2_from;
  edge2.from_len = 1;
  edge2_to[0] = env->node_c;
  edge2.to = edge2_to;
  edge2.to_len = 1;
  edge2.payload = env->payload2;

  memset(&edge3, 0, sizeof(edge3));
  edge3.type = 0x20;
  edge3_from[0] = env->node_b;
  edge3.from = edge3_from;
  edge3.from_len = 1;
  edge3_to[0] = env->node_a;
  edge3.to = edge3_to;
  edge3.to_len = 1;
  edge3.payload = env->payload3;

  memset(&edge_bad, 0, sizeof(edge_bad));
  edge_bad.type = 0x99;
  edge_bad_from[0] = env->node_a;
  edge_bad.from = edge_bad_from;
  edge_bad.from_len = 1;
  edge_bad_to[0] = env->node_b;
  edge_bad.to = edge_bad_to;
  edge_bad.to_len = 1;
  edge_bad.payload = env->payload_bad;

  if (!amduat_enc_tgk1_edge_encode_v1(&edge1, &env->edge_bytes[0]) ||
      !amduat_enc_tgk1_edge_encode_v1(&edge2, &env->edge_bytes[1]) ||
      !amduat_enc_tgk1_edge_encode_v1(&edge3, &env->edge_bytes[2]) ||
      !amduat_enc_tgk1_edge_encode_v1(&edge_bad, &env->edge_bytes[3])) {
    free_edge_bytes(env);
    return false;
  }

  env->artifacts[0].ref = env->ref_edge1;
  env->artifacts[0].artifact =
      amduat_artifact_with_type(env->edge_bytes[0],
                                amduat_type_tag(TYPE_TAG_TGK1_EDGE_V1));
  env->artifacts[1].ref = env->ref_edge2;
  env->artifacts[1].artifact =
      amduat_artifact_with_type(env->edge_bytes[1],
                                amduat_type_tag(TYPE_TAG_TGK1_EDGE_V1));
  env->artifacts[2].ref = env->ref_edge3;
  env->artifacts[2].artifact =
      amduat_artifact_with_type(env->edge_bytes[2],
                                amduat_type_tag(TYPE_TAG_TGK1_EDGE_V1));
  env->artifacts[3].ref = env->ref_edge_bad;
  env->artifacts[3].artifact =
      amduat_artifact_with_type(env->edge_bytes[3],
                                amduat_type_tag(TYPE_TAG_TGK1_EDGE_V1));
  env->artifacts[4].ref = env->ref_non_edge;
  env->artifacts[4].artifact =
      amduat_artifact(amduat_octets(k_non_edge_bytes, sizeof(k_non_edge_bytes)));

  if (!amduat_tgk_store_mem_init(&env->mem, env->config, env->artifacts, 5)) {
    free_edge_bytes(env);
    return false;
  }

  amduat_tgk_store_init(&env->store, env->config, amduat_tgk_store_mem_ops(),
                        &env->mem);
  return true;
}

static void free_env(test_env_t *env) {
  amduat_tgk_store_mem_free(&env->mem);
  free_edge_bytes(env);
}

static int test_resolve_edge_ok(const test_env_t *env) {
  amduat_tgk_edge_body_t body;
  amduat_tgk_graph_error_t err;

  err = amduat_tgk_store_resolve_edge((amduat_tgk_store_t *)&env->store,
                                      env->ref_edge1, &body);
  if (err != 0) {
    fprintf(stderr, "resolve_edge ok failed: %d\n", err);
    return 1;
  }
  if (body.type != 0x10 || body.from_len != 1 || body.to_len != 1) {
    fprintf(stderr, "resolve_edge body mismatch\n");
    amduat_tgk_edge_body_free(&body);
    return 1;
  }
  if (!amduat_reference_eq(body.from[0], env->node_a) ||
      !amduat_reference_eq(body.to[0], env->node_b)) {
    fprintf(stderr, "resolve_edge endpoints mismatch\n");
    amduat_tgk_edge_body_free(&body);
    return 1;
  }
  amduat_tgk_edge_body_free(&body);
  return 0;
}

static int test_resolve_edge_not_edge(const test_env_t *env) {
  amduat_tgk_edge_body_t body;
  amduat_tgk_graph_error_t err;

  err = amduat_tgk_store_resolve_edge((amduat_tgk_store_t *)&env->store,
                                      env->ref_edge_bad, &body);
  if (err != GS_ERR_NOT_EDGE) {
    fprintf(stderr, "resolve_edge not_edge mismatch: %d\n", err);
    return 1;
  }
  return 0;
}

static int test_resolve_edge_missing(const test_env_t *env) {
  amduat_tgk_edge_body_t body;
  amduat_tgk_graph_error_t err;

  err = amduat_tgk_store_resolve_edge((amduat_tgk_store_t *)&env->store,
                                      env->ref_missing, &body);
  if (err != GS_ERR_ARTIFACT_ERROR) {
    fprintf(stderr, "resolve_edge missing mismatch: %d\n", err);
    return 1;
  }
  return 0;
}

static bool make_ref_for_hash(amduat_hash_id_t hash_id,
                              uint8_t value,
                              amduat_reference_t *out_ref) {
  const amduat_hash_asl1_desc_t *desc;
  uint8_t *digest;

  if (out_ref == NULL) {
    return false;
  }
  desc = amduat_hash_asl1_desc_lookup(hash_id);
  if (desc == NULL || desc->digest_len == 0) {
    return false;
  }
  digest = (uint8_t *)malloc(desc->digest_len);
  if (digest == NULL) {
    return false;
  }
  memset(digest, value, desc->digest_len);
  *out_ref = amduat_reference(hash_id, amduat_octets(digest, desc->digest_len));
  return true;
}

static int test_resolve_edge_unsupported(const test_env_t *env) {
  amduat_reference_t ref;
  amduat_tgk_edge_body_t body;
  amduat_tgk_graph_error_t err;
  int exit_code = 1;

  if (!make_ref_for_hash(0x0002, 0x6au, &ref)) {
    fprintf(stderr, "unsupported ref alloc failed\n");
    return 1;
  }

  err = amduat_tgk_store_resolve_edge((amduat_tgk_store_t *)&env->store, ref,
                                      &body);
  if (err != GS_ERR_UNSUPPORTED) {
    fprintf(stderr, "resolve_edge unsupported mismatch: %d\n", err);
    goto cleanup;
  }

  exit_code = 0;
cleanup:
  free_ref(&ref);
  return exit_code;
}

static int test_type_filter(const test_env_t *env) {
  amduat_tgk_edge_type_id_t types[1] = {0x10};
  amduat_tgk_edge_type_filter_t filter;
  amduat_tgk_graph_edge_view_list_t edges;
  size_t i;

  filter.types = types;
  filter.types_len = 1;

  if (!amduat_tgk_store_edges_incident((amduat_tgk_store_t *)&env->store,
                                       env->node_a, filter, &edges)) {
    fprintf(stderr, "edges_incident type filter failed\n");
    return 1;
  }
  if (edges.len != 2) {
    fprintf(stderr, "edges_incident type filter count mismatch\n");
    amduat_tgk_graph_edge_view_list_free(&edges);
    return 1;
  }
  for (i = 0; i < edges.len; ++i) {
    if (edges.edges[i].body.type != 0x10) {
      fprintf(stderr, "edges_incident type filter type mismatch\n");
      amduat_tgk_graph_edge_view_list_free(&edges);
      return 1;
    }
  }
  amduat_tgk_graph_edge_view_list_free(&edges);
  return 0;
}

static int test_ordering(const test_env_t *env) {
  amduat_tgk_edge_type_filter_t filter;
  amduat_tgk_graph_edge_view_list_t edges;

  filter.types = NULL;
  filter.types_len = 0;

  if (!amduat_tgk_store_edges_from((amduat_tgk_store_t *)&env->store,
                                   env->node_a, filter, &edges)) {
    fprintf(stderr, "edges_from ordering failed\n");
    return 1;
  }
  if (edges.len != 2) {
    fprintf(stderr, "edges_from ordering count mismatch\n");
    amduat_tgk_graph_edge_view_list_free(&edges);
    return 1;
  }
  if (!amduat_reference_eq(edges.edges[0].edge_ref, env->ref_edge2) ||
      !amduat_reference_eq(edges.edges[1].edge_ref, env->ref_edge1)) {
    fprintf(stderr, "edges_from ordering mismatch\n");
    amduat_tgk_graph_edge_view_list_free(&edges);
    return 1;
  }
  amduat_tgk_graph_edge_view_list_free(&edges);
  return 0;
}

static int test_adjacency(const test_env_t *env) {
  amduat_tgk_edge_type_filter_t filter;
  amduat_tgk_graph_edge_view_list_t edges;

  filter.types = NULL;
  filter.types_len = 0;

  if (!amduat_tgk_store_edges_to((amduat_tgk_store_t *)&env->store,
                                 env->node_a, filter, &edges)) {
    fprintf(stderr, "edges_to failed\n");
    return 1;
  }
  if (edges.len != 1 ||
      !amduat_reference_eq(edges.edges[0].edge_ref, env->ref_edge3)) {
    fprintf(stderr, "edges_to mismatch\n");
    amduat_tgk_graph_edge_view_list_free(&edges);
    return 1;
  }
  amduat_tgk_graph_edge_view_list_free(&edges);

  if (!amduat_tgk_store_edges_incident((amduat_tgk_store_t *)&env->store,
                                       env->node_a, filter, &edges)) {
    fprintf(stderr, "edges_incident failed\n");
    return 1;
  }
  if (edges.len != 3 ||
      !amduat_reference_eq(edges.edges[0].edge_ref, env->ref_edge2) ||
      !amduat_reference_eq(edges.edges[1].edge_ref, env->ref_edge1) ||
      !amduat_reference_eq(edges.edges[2].edge_ref, env->ref_edge3)) {
    fprintf(stderr, "edges_incident mismatch\n");
    amduat_tgk_graph_edge_view_list_free(&edges);
    return 1;
  }
  amduat_tgk_graph_edge_view_list_free(&edges);

  return 0;
}

static int test_neighbors(const test_env_t *env) {
  amduat_tgk_edge_type_filter_t filter;
  amduat_tgk_node_list_t nodes;

  filter.types = NULL;
  filter.types_len = 0;

  if (!amduat_tgk_store_neighbors((amduat_tgk_store_t *)&env->store,
                                  env->node_a, filter,
                                  AMDUAT_TGK_GRAPH_DIR_OUT, &nodes)) {
    fprintf(stderr, "neighbors out failed\n");
    return 1;
  }
  if (nodes.len != 2 ||
      !amduat_reference_eq(nodes.nodes[0], env->node_b) ||
      !amduat_reference_eq(nodes.nodes[1], env->node_c)) {
    fprintf(stderr, "neighbors out mismatch\n");
    amduat_tgk_node_list_free(&nodes);
    return 1;
  }
  amduat_tgk_node_list_free(&nodes);

  if (!amduat_tgk_store_neighbors((amduat_tgk_store_t *)&env->store,
                                  env->node_a, filter,
                                  AMDUAT_TGK_GRAPH_DIR_IN, &nodes)) {
    fprintf(stderr, "neighbors in failed\n");
    return 1;
  }
  if (nodes.len != 1 ||
      !amduat_reference_eq(nodes.nodes[0], env->node_b)) {
    fprintf(stderr, "neighbors in mismatch\n");
    amduat_tgk_node_list_free(&nodes);
    return 1;
  }
  amduat_tgk_node_list_free(&nodes);

  if (!amduat_tgk_store_neighbors((amduat_tgk_store_t *)&env->store,
                                  env->node_a, filter,
                                  AMDUAT_TGK_GRAPH_DIR_BOTH, &nodes)) {
    fprintf(stderr, "neighbors both failed\n");
    return 1;
  }
  if (nodes.len != 2 ||
      !amduat_reference_eq(nodes.nodes[0], env->node_b) ||
      !amduat_reference_eq(nodes.nodes[1], env->node_c)) {
    fprintf(stderr, "neighbors both mismatch\n");
    amduat_tgk_node_list_free(&nodes);
    return 1;
  }
  amduat_tgk_node_list_free(&nodes);

  return 0;
}

static int test_scan_edges(const test_env_t *env) {
  amduat_tgk_edge_type_filter_t filter;
  amduat_tgk_graph_scan_result_t scan;

  filter.types = NULL;
  filter.types_len = 0;

  if (!amduat_tgk_store_scan_edges((amduat_tgk_store_t *)&env->store, filter,
                                   amduat_octets(NULL, 0), false, &scan)) {
    fprintf(stderr, "scan_edges failed\n");
    return 1;
  }
  if (scan.edges.len != 3 ||
      !amduat_reference_eq(scan.edges.edges[0].edge_ref, env->ref_edge2) ||
      !amduat_reference_eq(scan.edges.edges[1].edge_ref, env->ref_edge1) ||
      !amduat_reference_eq(scan.edges.edges[2].edge_ref, env->ref_edge3)) {
    fprintf(stderr, "scan_edges ordering mismatch\n");
    amduat_tgk_graph_scan_result_free(&scan);
    return 1;
  }
  amduat_tgk_graph_scan_result_free(&scan);
  return 0;
}

int main(void) {
  test_env_t env;

  if (!init_env(&env)) {
    fprintf(stderr, "failed to init tgk store mem test env\n");
    return 1;
  }

  if (test_resolve_edge_ok(&env) != 0 ||
      test_resolve_edge_not_edge(&env) != 0 ||
      test_resolve_edge_missing(&env) != 0 ||
      test_resolve_edge_unsupported(&env) != 0 ||
      test_type_filter(&env) != 0 ||
      test_ordering(&env) != 0 ||
      test_adjacency(&env) != 0 ||
      test_neighbors(&env) != 0 ||
      test_scan_edges(&env) != 0) {
    free_env(&env);
    return 1;
  }

  free_env(&env);
  return 0;
}