#include "amduat/pel/program_dag.h"
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

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 bytes_equal(amduat_octets_t bytes,
                        const uint8_t *expected,
                        size_t expected_len) {
  if (bytes.len != expected_len) {
    return false;
  }
  if (bytes.len == 0) {
    return true;
  }
  return memcmp(bytes.data, expected, expected_len) == 0;
}

static int test_valid_program(void) {
  amduat_pel_dag_input_t node2_inputs[2];
  amduat_pel_node_t nodes[2];
  amduat_pel_root_ref_t roots[1];
  amduat_pel_program_t program;
  amduat_artifact_t inputs[1];
  amduat_artifact_t *outputs = NULL;
  size_t outputs_len = 0;
  amduat_pel_execution_result_value_t result;
  const char op_const[] = "pel.bytes.const";
  const char op_concat[] = "pel.bytes.concat";
  const uint8_t const_bytes[] = {'h', 'i'};
  const uint8_t expected[] = {'h', 'i', '!'};
  amduat_octets_t const_params;
  int exit_code = 1;

  const_params = make_const_params(const_bytes, sizeof(const_bytes));
  if (const_params.data == NULL) {
    fprintf(stderr, "const params alloc failed\n");
    return exit_code;
  }

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

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

  node2_inputs[0].kind = AMDUAT_PEL_DAG_INPUT_NODE;
  node2_inputs[0].value.node.node_id = 1;
  node2_inputs[0].value.node.output_index = 0;
  node2_inputs[1].kind = AMDUAT_PEL_DAG_INPUT_EXTERNAL;
  node2_inputs[1].value.external.input_index = 0;

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

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

  inputs[0] = amduat_artifact(amduat_octets("!", 1));

  if (!amduat_pel_program_dag_exec(&program, inputs, 1, NULL, &outputs,
                                   &outputs_len, &result)) {
    fprintf(stderr, "exec failed\n");
    goto cleanup;
  }

  if (result.status != AMDUAT_PEL_EXEC_STATUS_OK ||
      result.summary.kind != AMDUAT_PEL_EXEC_ERROR_NONE ||
      result.summary.status_code != 0) {
    fprintf(stderr, "unexpected status\n");
    goto cleanup;
  }

  if (outputs_len != 1 ||
      !bytes_equal(outputs[0].bytes, expected, sizeof(expected))) {
    fprintf(stderr, "unexpected output\n");
    goto cleanup;
  }

  exit_code = 0;

cleanup:
  amduat_pel_program_dag_free_outputs(outputs, outputs_len);
  free((void *)const_params.data);
  return exit_code;
}

static int test_invalid_program_cycle(void) {
  amduat_pel_dag_input_t node1_inputs[1];
  amduat_pel_dag_input_t node2_inputs[1];
  amduat_pel_node_t nodes[2];
  amduat_pel_root_ref_t roots[1];
  amduat_pel_program_t program;
  amduat_artifact_t *outputs = NULL;
  size_t outputs_len = 0;
  amduat_pel_execution_result_value_t result;
  const char op_concat[] = "pel.bytes.concat";

  node1_inputs[0].kind = AMDUAT_PEL_DAG_INPUT_NODE;
  node1_inputs[0].value.node.node_id = 2;
  node1_inputs[0].value.node.output_index = 0;

  node2_inputs[0].kind = AMDUAT_PEL_DAG_INPUT_NODE;
  node2_inputs[0].value.node.node_id = 1;
  node2_inputs[0].value.node.output_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 = node1_inputs;
  nodes[0].inputs_len = 1;
  nodes[0].params = amduat_octets(NULL, 0);

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

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

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

  if (!amduat_pel_program_dag_exec(&program, NULL, 0, NULL, &outputs,
                                   &outputs_len, &result)) {
    fprintf(stderr, "exec failed\n");
    return 1;
  }

  if (result.status != AMDUAT_PEL_EXEC_STATUS_INVALID_PROGRAM ||
      result.summary.kind != AMDUAT_PEL_EXEC_ERROR_PROGRAM ||
      result.summary.status_code != 2) {
    fprintf(stderr, "unexpected invalid program status\n");
    amduat_pel_program_dag_free_outputs(outputs, outputs_len);
    return 1;
  }

  amduat_pel_program_dag_free_outputs(outputs, outputs_len);
  return 0;
}

static int test_invalid_params(void) {
  amduat_pel_dag_input_t node_inputs[1];
  amduat_pel_node_t nodes[1];
  amduat_pel_root_ref_t roots[1];
  amduat_pel_program_t program;
  amduat_artifact_t inputs[1];
  amduat_artifact_t *outputs = NULL;
  size_t outputs_len = 0;
  amduat_pel_execution_result_value_t result;
  const char op_concat[] = "pel.bytes.concat";
  uint8_t bad_params[] = {0x01u};

  node_inputs[0].kind = AMDUAT_PEL_DAG_INPUT_EXTERNAL;
  node_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 = node_inputs;
  nodes[0].inputs_len = 1;
  nodes[0].params = amduat_octets(bad_params, sizeof(bad_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;

  inputs[0] = amduat_artifact(amduat_octets("x", 1));

  if (!amduat_pel_program_dag_exec(&program, inputs, 1, NULL, &outputs,
                                   &outputs_len, &result)) {
    fprintf(stderr, "exec failed\n");
    return 1;
  }

  if (result.status != AMDUAT_PEL_EXEC_STATUS_INVALID_PROGRAM ||
      result.summary.kind != AMDUAT_PEL_EXEC_ERROR_PROGRAM ||
      result.summary.status_code != 2) {
    fprintf(stderr, "unexpected invalid params status\n");
    amduat_pel_program_dag_free_outputs(outputs, outputs_len);
    return 1;
  }

  amduat_pel_program_dag_free_outputs(outputs, outputs_len);
  return 0;
}

static int test_invalid_input_index(void) {
  amduat_pel_dag_input_t node_inputs[1];
  amduat_pel_node_t nodes[1];
  amduat_pel_root_ref_t roots[1];
  amduat_pel_program_t program;
  amduat_artifact_t inputs[1];
  amduat_artifact_t *outputs = NULL;
  size_t outputs_len = 0;
  amduat_pel_execution_result_value_t result;
  const char op_concat[] = "pel.bytes.concat";

  node_inputs[0].kind = AMDUAT_PEL_DAG_INPUT_EXTERNAL;
  node_inputs[0].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 = node_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;

  inputs[0] = amduat_artifact(amduat_octets("x", 1));

  if (!amduat_pel_program_dag_exec(&program, inputs, 1, NULL, &outputs,
                                   &outputs_len, &result)) {
    fprintf(stderr, "exec failed\n");
    return 1;
  }

  if (result.status != AMDUAT_PEL_EXEC_STATUS_INVALID_INPUTS ||
      result.summary.kind != AMDUAT_PEL_EXEC_ERROR_INPUTS ||
      result.summary.status_code != 3) {
    fprintf(stderr, "unexpected invalid inputs status\n");
    amduat_pel_program_dag_free_outputs(outputs, outputs_len);
    return 1;
  }

  amduat_pel_program_dag_free_outputs(outputs, outputs_len);
  return 0;
}

int main(void) {
  if (test_valid_program() != 0) {
    return 1;
  }
  if (test_invalid_program_cycle() != 0) {
    return 1;
  }
  if (test_invalid_params() != 0) {
    return 1;
  }
  if (test_invalid_input_index() != 0) {
    return 1;
  }
  return 0;
}