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

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

typedef struct {
  amduat_tgk_store_t store;
  amduat_tgk_store_mem_t mem;
  amduat_tgk_store_mem_artifact_t artifacts[4];
  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_edge4;
  amduat_reference_t node_a;
  amduat_reference_t node_b;
  amduat_reference_t node_c;
  amduat_reference_t node_d;
  amduat_reference_t payload1;
  amduat_reference_t payload2;
  amduat_reference_t payload3;
  amduat_reference_t payload4;
  uint8_t digest_edge1[32];
  uint8_t digest_edge2[32];
  uint8_t digest_edge3[32];
  uint8_t digest_edge4[32];
  uint8_t digest_node_a[32];
  uint8_t digest_node_b[32];
  uint8_t digest_node_c[32];
  uint8_t digest_node_d[32];
  uint8_t digest_payload1[32];
  uint8_t digest_payload2[32];
  uint8_t digest_payload3[32];
  uint8_t digest_payload4[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 edge4;
  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 edge4_from[1];
  amduat_reference_t edge4_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(0x11, env->digest_edge1);
  env->ref_edge2 = make_ref(0x12, env->digest_edge2);
  env->ref_edge3 = make_ref(0x13, env->digest_edge3);
  env->ref_edge4 = make_ref(0x14, env->digest_edge4);

  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->node_d = make_ref(0xd1, env->digest_node_d);

  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->payload4 = make_ref(0xe4, env->digest_payload4);

  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_b;
  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 = 0x10;
  edge3_from[0] = env->node_c;
  edge3.from = edge3_from;
  edge3.from_len = 1;
  edge3_to[0] = env->node_d;
  edge3.to = edge3_to;
  edge3.to_len = 1;
  edge3.payload = env->payload3;

  memset(&edge4, 0, sizeof(edge4));
  edge4.type = 0x20;
  edge4_from[0] = env->node_b;
  edge4.from = edge4_from;
  edge4.from_len = 1;
  edge4_to[0] = env->node_a;
  edge4.to = edge4_to;
  edge4.to_len = 1;
  edge4.payload = env->payload4;

  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(&edge4, &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_edge4;
  env->artifacts[3].artifact =
      amduat_artifact_with_type(env->edge_bytes[3],
                                amduat_type_tag(TYPE_TAG_TGK1_EDGE_V1));

  if (!amduat_tgk_store_mem_init(&env->mem, env->config, env->artifacts, 4)) {
    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 bool list_contains(const amduat_reference_t *nodes,
                           size_t len,
                           amduat_reference_t node) {
  size_t i;

  for (i = 0; i < len; ++i) {
    if (amduat_reference_eq(nodes[i], node)) {
      return true;
    }
  }
  return false;
}

static bool depth_map_get(const amduat_tgk_prov_depth_map_t *map,
                          amduat_reference_t node,
                          uint32_t *out_depth) {
  size_t i;

  for (i = 0; i < map->len; ++i) {
    if (amduat_reference_eq(map->entries[i].node, node)) {
      if (out_depth != NULL) {
        *out_depth = map->entries[i].depth;
      }
      return true;
    }
  }
  return false;
}

static int test_direction(const test_env_t *env) {
  amduat_tgk_prov_query_t query;
  amduat_tgk_node_list_t nodes;
  amduat_tgk_edge_type_id_t types[1] = {0x10};
  amduat_reference_t seeds[1];

  seeds[0] = env->node_b;
  query.type_filter.types = types;
  query.type_filter.types_len = 1;
  query.has_depth_limit = false;

  query.direction = AMDUAT_TGK_PROV_DIR_FORWARD;
  if (!amduat_tgk_prov_closure_nodes((amduat_tgk_store_t *)&env->store, seeds, 1,
                                     query, &nodes)) {
    fprintf(stderr, "direction forward closure failed\n");
    return 1;
  }
  if (nodes.len != 3 || !list_contains(nodes.nodes, nodes.len, env->node_b) ||
      !list_contains(nodes.nodes, nodes.len, env->node_c) ||
      !list_contains(nodes.nodes, nodes.len, env->node_d)) {
    fprintf(stderr, "direction forward closure mismatch\n");
    amduat_tgk_node_list_free(&nodes);
    return 1;
  }
  amduat_tgk_node_list_free(&nodes);

  query.direction = AMDUAT_TGK_PROV_DIR_BACKWARD;
  if (!amduat_tgk_prov_closure_nodes((amduat_tgk_store_t *)&env->store, seeds, 1,
                                     query, &nodes)) {
    fprintf(stderr, "direction backward closure failed\n");
    return 1;
  }
  if (nodes.len != 2 || !list_contains(nodes.nodes, nodes.len, env->node_a) ||
      !list_contains(nodes.nodes, nodes.len, env->node_b)) {
    fprintf(stderr, "direction backward closure mismatch\n");
    amduat_tgk_node_list_free(&nodes);
    return 1;
  }
  amduat_tgk_node_list_free(&nodes);

  query.direction = AMDUAT_TGK_PROV_DIR_BOTH;
  if (!amduat_tgk_prov_closure_nodes((amduat_tgk_store_t *)&env->store, seeds, 1,
                                     query, &nodes)) {
    fprintf(stderr, "direction both closure failed\n");
    return 1;
  }
  if (nodes.len != 4 || !list_contains(nodes.nodes, nodes.len, env->node_a) ||
      !list_contains(nodes.nodes, nodes.len, env->node_b) ||
      !list_contains(nodes.nodes, nodes.len, env->node_c) ||
      !list_contains(nodes.nodes, nodes.len, env->node_d)) {
    fprintf(stderr, "direction both closure mismatch\n");
    amduat_tgk_node_list_free(&nodes);
    return 1;
  }
  amduat_tgk_node_list_free(&nodes);

  return 0;
}

static int test_depth_limits(const test_env_t *env) {
  amduat_tgk_prov_query_t query;
  amduat_tgk_prov_depth_map_t depth_map;
  amduat_tgk_prov_layering_t layers;
  amduat_tgk_edge_type_id_t types[1] = {0x10};
  amduat_reference_t seeds[1];
  uint32_t depth = 0;

  seeds[0] = env->node_a;
  query.direction = AMDUAT_TGK_PROV_DIR_FORWARD;
  query.type_filter.types = types;
  query.type_filter.types_len = 1;
  query.has_depth_limit = false;

  if (!amduat_tgk_prov_depths((amduat_tgk_store_t *)&env->store, seeds, 1, query,
                              &depth_map)) {
    fprintf(stderr, "depth map unbounded failed\n");
    return 1;
  }
  if (depth_map.len != 4 ||
      !depth_map_get(&depth_map, env->node_a, &depth) || depth != 0 ||
      !depth_map_get(&depth_map, env->node_b, &depth) || depth != 1 ||
      !depth_map_get(&depth_map, env->node_c, &depth) || depth != 2 ||
      !depth_map_get(&depth_map, env->node_d, &depth) || depth != 3) {
    fprintf(stderr, "depth map unbounded mismatch\n");
    amduat_tgk_prov_depth_map_free(&depth_map);
    return 1;
  }
  amduat_tgk_prov_depth_map_free(&depth_map);

  query.has_depth_limit = true;
  query.depth_limit = 1;
  if (!amduat_tgk_prov_layers((amduat_tgk_store_t *)&env->store, seeds, 1, query,
                              &layers)) {
    fprintf(stderr, "layers bounded failed\n");
    return 1;
  }
  if (layers.layers_len != 2 || layers.layers[0].depth != 0 ||
      layers.layers[1].depth != 1 || layers.layers[0].nodes_len != 1 ||
      layers.layers[1].nodes_len != 1 ||
      !list_contains(layers.layers[0].nodes, layers.layers[0].nodes_len,
                     env->node_a) ||
      !list_contains(layers.layers[1].nodes, layers.layers[1].nodes_len,
                     env->node_b)) {
    fprintf(stderr, "layers bounded mismatch\n");
    amduat_tgk_prov_layering_free(&layers);
    return 1;
  }
  amduat_tgk_prov_layering_free(&layers);
  return 0;
}

static int test_payload_in_trace(const test_env_t *env) {
  amduat_tgk_prov_query_t query;
  amduat_tgk_trace_graph_t trace;
  amduat_tgk_node_list_t nodes;
  amduat_tgk_edge_type_id_t types[1] = {0x10};
  amduat_reference_t seeds[1];

  seeds[0] = env->node_a;
  query.direction = AMDUAT_TGK_PROV_DIR_FORWARD;
  query.type_filter.types = types;
  query.type_filter.types_len = 1;
  query.has_depth_limit = false;

  if (!amduat_tgk_prov_trace((amduat_tgk_store_t *)&env->store, seeds, 1, query,
                             &trace)) {
    fprintf(stderr, "trace failed\n");
    return 1;
  }
  if (trace.edges_len != 3 ||
      !list_contains(trace.nodes, trace.nodes_len, env->payload1) ||
      !list_contains(trace.nodes, trace.nodes_len, env->payload2) ||
      !list_contains(trace.nodes, trace.nodes_len, env->payload3)) {
    fprintf(stderr, "trace payload nodes mismatch\n");
    amduat_tgk_trace_graph_free(&trace);
    return 1;
  }
  amduat_tgk_trace_graph_free(&trace);

  if (!amduat_tgk_prov_closure_nodes((amduat_tgk_store_t *)&env->store, seeds, 1,
                                     query, &nodes)) {
    fprintf(stderr, "closure for payload test failed\n");
    return 1;
  }
  if (list_contains(nodes.nodes, nodes.len, env->payload1) ||
      list_contains(nodes.nodes, nodes.len, env->payload2) ||
      list_contains(nodes.nodes, nodes.len, env->payload3)) {
    fprintf(stderr, "closure should not include payload nodes\n");
    amduat_tgk_node_list_free(&nodes);
    return 1;
  }
  amduat_tgk_node_list_free(&nodes);

  return 0;
}

static int test_empty_seed(const test_env_t *env) {
  amduat_tgk_prov_query_t query;
  amduat_tgk_node_list_t nodes;
  amduat_tgk_prov_depth_map_t depth_map;
  amduat_tgk_prov_layering_t layers;
  amduat_tgk_trace_graph_t trace;
  amduat_tgk_edge_type_id_t types[1] = {0x10};

  query.direction = AMDUAT_TGK_PROV_DIR_FORWARD;
  query.type_filter.types = types;
  query.type_filter.types_len = 1;
  query.has_depth_limit = false;

  if (!amduat_tgk_prov_closure_nodes((amduat_tgk_store_t *)&env->store, NULL, 0,
                                     query, &nodes)) {
    fprintf(stderr, "empty seed closure failed\n");
    return 1;
  }
  if (nodes.len != 0) {
    fprintf(stderr, "empty seed closure not empty\n");
    amduat_tgk_node_list_free(&nodes);
    return 1;
  }
  amduat_tgk_node_list_free(&nodes);

  if (!amduat_tgk_prov_depths((amduat_tgk_store_t *)&env->store, NULL, 0, query,
                              &depth_map)) {
    fprintf(stderr, "empty seed depths failed\n");
    return 1;
  }
  if (depth_map.len != 0) {
    fprintf(stderr, "empty seed depths not empty\n");
    amduat_tgk_prov_depth_map_free(&depth_map);
    return 1;
  }
  amduat_tgk_prov_depth_map_free(&depth_map);

  if (!amduat_tgk_prov_layers((amduat_tgk_store_t *)&env->store, NULL, 0, query,
                              &layers)) {
    fprintf(stderr, "empty seed layers failed\n");
    return 1;
  }
  if (layers.layers_len != 0) {
    fprintf(stderr, "empty seed layers not empty\n");
    amduat_tgk_prov_layering_free(&layers);
    return 1;
  }
  amduat_tgk_prov_layering_free(&layers);

  if (!amduat_tgk_prov_trace((amduat_tgk_store_t *)&env->store, NULL, 0, query,
                             &trace)) {
    fprintf(stderr, "empty seed trace failed\n");
    return 1;
  }
  if (trace.seeds_len != 0 || trace.nodes_len != 0 || trace.edges_len != 0) {
    fprintf(stderr, "empty seed trace not empty\n");
    amduat_tgk_trace_graph_free(&trace);
    return 1;
  }
  amduat_tgk_trace_graph_free(&trace);

  return 0;
}

int main(void) {
  test_env_t env;

  if (!init_env(&env)) {
    fprintf(stderr, "failed to initialize test environment\n");
    return 1;
  }

  if (test_direction(&env) != 0 || test_depth_limits(&env) != 0 ||
      test_payload_in_trace(&env) != 0 || test_empty_seed(&env) != 0) {
    free_env(&env);
    return 1;
  }

  free_env(&env);
  return 0;
}