#include "amduat/pel/surf.h"

#include "amduat/enc/pel1_result.h"
#include "amduat/enc/pel_program_dag.h"
#include "amduat/enc/pel_trace_dag.h"
#include "pel_program_dag_internal.h"
#include "amduat/pel/program_dag_desc.h"

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

void amduat_pel_surf_free_refs(amduat_reference_t *refs, size_t refs_len) {
  size_t i;

  if (refs == NULL) {
    return;
  }
  for (i = 0; i < refs_len; ++i) {
    amduat_reference_free(&refs[i]);
  }
  free(refs);
}

void amduat_pel_surf_free_ref(amduat_reference_t *ref) {
  amduat_reference_free(ref);
}

static void amduat_init_core_result(amduat_pel_execution_result_value_t *result,
                                    amduat_reference_t scheme_ref,
                                    amduat_pel_execution_status_t status,
                                    amduat_pel_execution_error_kind_t kind,
                                    uint32_t status_code) {
  if (result == NULL) {
    return;
  }
  result->pel1_version = 1;
  result->status = status;
  result->scheme_ref = scheme_ref;
  result->summary.kind = kind;
  result->summary.status_code = status_code;
  result->diagnostics = NULL;
  result->diagnostics_len = 0;
}

static bool amduat_surf_diag_setf(
    amduat_pel_execution_result_value_t *result,
    uint32_t code,
    const char *fmt,
    ...) {
  va_list ap;
  va_list ap_copy;
  int needed;
  size_t message_len;
  uint8_t *buffer;
  amduat_pel_diagnostic_entry_t *entries;

  if (result == NULL || fmt == NULL) {
    return false;
  }

  amduat_pel_execution_result_free(result);

  va_start(ap, fmt);
  va_copy(ap_copy, ap);
  needed = vsnprintf(NULL, 0, fmt, ap_copy);
  va_end(ap_copy);
  if (needed < 0) {
    va_end(ap);
    return false;
  }

  message_len = (size_t)needed;
  buffer = NULL;
  if (message_len != 0) {
    buffer = (uint8_t *)malloc(message_len + 1);
    if (buffer == NULL) {
      va_end(ap);
      return false;
    }
    if (vsnprintf((char *)buffer, message_len + 1, fmt, ap) < 0) {
      free(buffer);
      va_end(ap);
      return false;
    }
  }
  va_end(ap);

  entries = (amduat_pel_diagnostic_entry_t *)calloc(1, sizeof(*entries));
  if (entries == NULL) {
    free(buffer);
    return false;
  }
  entries[0].code = code;
  entries[0].message = amduat_octets(buffer, message_len);
  result->diagnostics = entries;
  result->diagnostics_len = 1;
  return true;
}

enum {
  AMDUAT_PEL_DAG_DIAG_DECODE_FAILED = 0x00010001u,
  AMDUAT_PEL_DAG_DIAG_WRONG_TYPE_TAG = 0x00010002u
};

static bool amduat_store_error_map(amduat_asl_store_error_t err,
                                   amduat_pel_store_error_code_t *out_code) {
  if (out_code == NULL) {
    return false;
  }
  switch (err) {
    case AMDUAT_ASL_STORE_ERR_NOT_FOUND:
      *out_code = AMDUAT_PEL_STORE_ERROR_NOT_FOUND;
      return true;
    case AMDUAT_ASL_STORE_ERR_INTEGRITY:
      *out_code = AMDUAT_PEL_STORE_ERROR_INTEGRITY;
      return true;
    case AMDUAT_ASL_STORE_ERR_UNSUPPORTED:
      *out_code = AMDUAT_PEL_STORE_ERROR_UNSUPPORTED;
      return true;
    default:
      return false;
  }
}

static bool amduat_store_surface_result(
    amduat_asl_store_t *store,
    const amduat_pel_execution_result_value_t *core_result,
    amduat_reference_t scheme_ref,
    amduat_reference_t program_ref,
    const amduat_reference_t *input_refs,
    size_t input_refs_len,
    const amduat_reference_t *output_refs,
    size_t output_refs_len,
    bool has_params_ref,
    amduat_reference_t params_ref,
    bool has_store_failure,
    const amduat_pel_store_failure_t *store_failure,
    bool has_trace_ref,
    amduat_reference_t trace_ref,
    amduat_reference_t *out_result_ref) {
  amduat_pel_surface_execution_result_t result;
  amduat_octets_t encoded;
  amduat_artifact_t artifact;
  amduat_asl_store_error_t store_err;

  if (store == NULL || core_result == NULL || out_result_ref == NULL) {
    return false;
  }

  memset(&result, 0, sizeof(result));
  result.pel1_version = 1;
  result.core_result = *core_result;
  result.scheme_ref = scheme_ref;
  result.program_ref = program_ref;
  result.input_refs = (amduat_reference_t *)input_refs;
  result.input_refs_len = input_refs_len;
  result.output_refs = (amduat_reference_t *)output_refs;
  result.output_refs_len = output_refs_len;
  result.has_params_ref = has_params_ref;
  if (has_params_ref) {
    result.params_ref = params_ref;
  }
  result.has_store_failure = has_store_failure;
  if (has_store_failure && store_failure != NULL) {
    result.store_failure = *store_failure;
  }
  result.has_trace_ref = has_trace_ref;
  if (has_trace_ref) {
    result.trace_ref = trace_ref;
  }

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

  artifact = amduat_artifact_with_type(
      encoded, amduat_type_tag(AMDUAT_TYPE_TAG_PEL1_RESULT_1));
  store_err = amduat_asl_store_put(store, artifact, out_result_ref);
  free((void *)encoded.data);
  return store_err == AMDUAT_ASL_STORE_OK;
}

static void amduat_trace_nodes_free(amduat_pel_node_trace_dag_t *nodes,
                                    size_t nodes_len) {
  size_t i;

  if (nodes == NULL) {
    return;
  }
  for (i = 0; i < nodes_len; ++i) {
    amduat_pel_node_trace_dag_t *node = &nodes[i];
    if (node->output_refs != NULL) {
      size_t j;
      for (j = 0; j < node->output_refs_len; ++j) {
        amduat_reference_free(&node->output_refs[j]);
      }
      free(node->output_refs);
      node->output_refs = NULL;
      node->output_refs_len = 0;
    }
  }
  free(nodes);
}

static bool amduat_store_trace(
    amduat_asl_store_t *store,
    amduat_reference_t scheme_ref,
    amduat_reference_t program_ref,
    const amduat_reference_t *input_refs,
    size_t input_refs_len,
    bool has_params_ref,
    amduat_reference_t params_ref,
    const amduat_pel_execution_result_value_t *core_result,
    const amduat_pel_program_t *program,
    const amduat_pel_program_dag_trace_t *trace_eval,
    amduat_reference_t *out_trace_ref) {
  amduat_pel_trace_dag_value_t trace;
  amduat_pel_node_trace_dag_t *node_traces;
  amduat_octets_t encoded;
  amduat_artifact_t artifact;
  amduat_asl_store_error_t store_err;
  size_t node_count;
  size_t i;

  if (store == NULL || core_result == NULL || out_trace_ref == NULL) {
    return false;
  }

  node_traces = NULL;
  node_count = 0;
  if (program != NULL && trace_eval != NULL && trace_eval->any_node_executed) {
    node_count = program->nodes_len;
    if (node_count != 0) {
      node_traces = (amduat_pel_node_trace_dag_t *)calloc(
          node_count, sizeof(*node_traces));
      if (node_traces == NULL) {
        return false;
      }
    }

    for (i = 0; i < node_count; ++i) {
      size_t node_index = trace_eval->order[i];
      const amduat_pel_node_t *node = &program->nodes[node_index];
      amduat_pel_program_dag_node_result_t *state =
          &trace_eval->nodes[node_index];
      amduat_pel_node_trace_dag_t *trace_node = &node_traces[i];

      trace_node->node_id = node->id;
      trace_node->op_name = node->op.name;
      trace_node->op_version = node->op.version;
      trace_node->status = state->status;
      trace_node->status_code = state->status_code;
      trace_node->diagnostics = NULL;
      trace_node->diagnostics_len = 0;

      if (state->status == AMDUAT_PEL_NODE_TRACE_OK &&
          state->outputs_len != 0) {
        size_t j;
        trace_node->output_refs = (amduat_reference_t *)calloc(
            state->outputs_len, sizeof(*trace_node->output_refs));
        if (trace_node->output_refs == NULL) {
          amduat_trace_nodes_free(node_traces, node_count);
          return false;
        }
        trace_node->output_refs_len = state->outputs_len;
        for (j = 0; j < state->outputs_len; ++j) {
          store_err = amduat_asl_store_put(store, state->outputs[j],
                                           &trace_node->output_refs[j]);
          if (store_err != AMDUAT_ASL_STORE_OK) {
            amduat_trace_nodes_free(node_traces, node_count);
            return false;
          }
        }
      }
    }
  }

  memset(&trace, 0, sizeof(trace));
  trace.pel1_version = 1;
  trace.scheme_ref = scheme_ref;
  trace.program_ref = program_ref;
  trace.status = core_result->status;
  trace.summary = core_result->summary;
  trace.has_exec_result_ref = false;
  trace.input_refs = (amduat_reference_t *)input_refs;
  trace.input_refs_len = input_refs_len;
  trace.has_params_ref = has_params_ref;
  if (has_params_ref) {
    trace.params_ref = params_ref;
  }
  trace.node_traces = node_traces;
  trace.node_traces_len = node_count;

  encoded = amduat_octets(NULL, 0);
  if (!amduat_enc_pel_trace_dag_encode_v1(&trace, &encoded)) {
    amduat_trace_nodes_free(node_traces, node_count);
    return false;
  }

  artifact = amduat_artifact_with_type(
      encoded, amduat_type_tag(AMDUAT_TYPE_TAG_PEL_TRACE_DAG_1));
  store_err = amduat_asl_store_put(store, artifact, out_trace_ref);
  free((void *)encoded.data);
  amduat_trace_nodes_free(node_traces, node_count);
  return store_err == AMDUAT_ASL_STORE_OK;
}

bool amduat_pel_surf_run(amduat_asl_store_t *store,
                         amduat_reference_t scheme_ref,
                         amduat_reference_t program_ref,
                         const amduat_reference_t *input_refs,
                         size_t input_refs_len,
                         bool has_params_ref,
                         amduat_reference_t params_ref,
                         amduat_reference_t **out_output_refs,
                         size_t *out_output_refs_len,
                         amduat_reference_t *out_result_ref) {
  amduat_pel_execution_result_value_t core_result;
  amduat_pel_store_failure_t store_failure;
  amduat_pel_store_error_code_t store_error_code;
  amduat_asl_store_error_t store_err;
  amduat_artifact_t program_artifact;
  amduat_artifact_t *input_artifacts;
  amduat_artifact_t params_artifact;
  amduat_pel_program_t program;
  amduat_pel_program_dag_trace_t trace_eval;
  amduat_artifact_t *outputs;
  size_t outputs_len;
  amduat_reference_t *output_refs;
  size_t output_refs_len;
  amduat_reference_t trace_ref;
  size_t i;
  bool program_decoded = false;
  bool has_params_artifact = false;
  bool exec_invoked = false;
  bool trace_ok = false;

  if (store == NULL || out_output_refs == NULL ||
      out_output_refs_len == NULL || out_result_ref == NULL) {
    return false;
  }
  if (input_refs_len != 0 && input_refs == NULL) {
    return false;
  }

  *out_output_refs = NULL;
  *out_output_refs_len = 0;
  out_result_ref->hash_id = 0;
  out_result_ref->digest = amduat_octets(NULL, 0);
  trace_ref.hash_id = 0;
  trace_ref.digest = amduat_octets(NULL, 0);

  if (!amduat_reference_eq(scheme_ref,
                           amduat_pel_program_dag_scheme_ref())) {
    amduat_init_core_result(&core_result, scheme_ref,
                            AMDUAT_PEL_EXEC_STATUS_SCHEME_UNSUPPORTED,
                            AMDUAT_PEL_EXEC_ERROR_SCHEME, 1);
    return amduat_store_surface_result(store, &core_result, scheme_ref,
                                       program_ref, input_refs, input_refs_len,
                                       NULL, 0, has_params_ref, params_ref,
                                       false, NULL, false, trace_ref,
                                       out_result_ref);
  }

  store_err = amduat_asl_store_get(store, program_ref, &program_artifact);
  if (store_err != AMDUAT_ASL_STORE_OK) {
    if (!amduat_store_error_map(store_err, &store_error_code)) {
      return false;
    }
    amduat_init_core_result(&core_result, scheme_ref,
                            AMDUAT_PEL_EXEC_STATUS_INVALID_PROGRAM,
                            AMDUAT_PEL_EXEC_ERROR_PROGRAM, 2);
    store_failure.phase = AMDUAT_PEL_STORE_FAILURE_PROGRAM;
    store_failure.error_code = store_error_code;
    store_failure.failing_ref = program_ref;
    return amduat_store_surface_result(
        store, &core_result, scheme_ref, program_ref, input_refs,
        input_refs_len, NULL, 0, has_params_ref, params_ref, true,
        &store_failure, false, trace_ref, out_result_ref);
  }

  input_artifacts = NULL;
  if (input_refs_len != 0) {
    input_artifacts = (amduat_artifact_t *)calloc(
        input_refs_len, sizeof(*input_artifacts));
    if (input_artifacts == NULL) {
      amduat_artifact_free(&program_artifact);
      return false;
    }
  }

  for (i = 0; i < input_refs_len; ++i) {
    store_err = amduat_asl_store_get(store, input_refs[i],
                                     &input_artifacts[i]);
    if (store_err != AMDUAT_ASL_STORE_OK) {
      size_t j;
      if (!amduat_store_error_map(store_err, &store_error_code)) {
        amduat_artifact_free(&program_artifact);
        free(input_artifacts);
        return false;
      }
      for (j = 0; j < i; ++j) {
        amduat_artifact_free(&input_artifacts[j]);
      }
      free(input_artifacts);
      amduat_artifact_free(&program_artifact);
      amduat_init_core_result(&core_result, scheme_ref,
                              AMDUAT_PEL_EXEC_STATUS_INVALID_INPUTS,
                              AMDUAT_PEL_EXEC_ERROR_INPUTS, 3);
      store_failure.phase = AMDUAT_PEL_STORE_FAILURE_INPUT;
      store_failure.error_code = store_error_code;
      store_failure.failing_ref = input_refs[i];
      return amduat_store_surface_result(
          store, &core_result, scheme_ref, program_ref, input_refs,
          input_refs_len, NULL, 0, has_params_ref, params_ref, true,
          &store_failure, false, trace_ref, out_result_ref);
    }
  }

  if (has_params_ref) {
    store_err = amduat_asl_store_get(store, params_ref, &params_artifact);
    if (store_err != AMDUAT_ASL_STORE_OK) {
      if (!amduat_store_error_map(store_err, &store_error_code)) {
        amduat_artifact_free(&program_artifact);
        for (i = 0; i < input_refs_len; ++i) {
          amduat_artifact_free(&input_artifacts[i]);
        }
        free(input_artifacts);
        return false;
      }
      amduat_artifact_free(&program_artifact);
      for (i = 0; i < input_refs_len; ++i) {
        amduat_artifact_free(&input_artifacts[i]);
      }
      free(input_artifacts);
      amduat_init_core_result(&core_result, scheme_ref,
                              AMDUAT_PEL_EXEC_STATUS_INVALID_INPUTS,
                              AMDUAT_PEL_EXEC_ERROR_INPUTS, 3);
      store_failure.phase = AMDUAT_PEL_STORE_FAILURE_INPUT;
      store_failure.error_code = store_error_code;
      store_failure.failing_ref = params_ref;
      return amduat_store_surface_result(
          store, &core_result, scheme_ref, program_ref, input_refs,
          input_refs_len, NULL, 0, has_params_ref, params_ref, true,
          &store_failure, false, trace_ref, out_result_ref);
    }
    has_params_artifact = true;
  }

  exec_invoked = true;
  outputs = NULL;
  outputs_len = 0;
  memset(&trace_eval, 0, sizeof(trace_eval));
  if (!program_artifact.has_type_tag ||
      program_artifact.type_tag.tag_id !=
          AMDUAT_PEL_TYPE_TAG_PROGRAM_DAG_1) {
    amduat_init_core_result(&core_result, scheme_ref,
                            AMDUAT_PEL_EXEC_STATUS_INVALID_PROGRAM,
                            AMDUAT_PEL_EXEC_ERROR_PROGRAM, 2);
    amduat_surf_diag_setf(&core_result,
                          AMDUAT_PEL_DAG_DIAG_WRONG_TYPE_TAG,
                          "invalid program: wrong type tag");
  } else {
    amduat_pel_program_dag_decode_status_t decode_status =
        amduat_enc_pel_program_dag_decode_v1_ex(program_artifact.bytes,
                                                &program);
    if (decode_status == AMDUAT_PEL_PROGRAM_DAG_DECODE_OOM) {
      amduat_artifact_free(&program_artifact);
      for (i = 0; i < input_refs_len; ++i) {
        amduat_artifact_free(&input_artifacts[i]);
      }
      free(input_artifacts);
      if (has_params_artifact) {
        amduat_artifact_free(&params_artifact);
      }
      return false;
    }
    if (decode_status != AMDUAT_PEL_PROGRAM_DAG_DECODE_OK) {
      amduat_init_core_result(&core_result, scheme_ref,
                              AMDUAT_PEL_EXEC_STATUS_INVALID_PROGRAM,
                              AMDUAT_PEL_EXEC_ERROR_PROGRAM, 2);
      amduat_surf_diag_setf(&core_result,
                            AMDUAT_PEL_DAG_DIAG_DECODE_FAILED,
                            "invalid program: decode failed");
    } else {
      program_decoded = true;
      const amduat_artifact_t *params_arg =
          has_params_artifact ? &params_artifact : NULL;
      if (!amduat_pel_program_dag_exec_trace(
              &program, input_artifacts, input_refs_len, params_arg, &outputs,
              &outputs_len, &core_result, &trace_eval)) {
        amduat_artifact_free(&program_artifact);
        amduat_enc_pel_program_dag_free(&program);
        for (i = 0; i < input_refs_len; ++i) {
          amduat_artifact_free(&input_artifacts[i]);
        }
        free(input_artifacts);
        if (has_params_artifact) {
          amduat_artifact_free(&params_artifact);
        }
        amduat_pel_program_dag_free_outputs(outputs, outputs_len);
        return false;
      }
    }
  }

  amduat_artifact_free(&program_artifact);

  output_refs = NULL;
  output_refs_len = outputs_len;
  if (outputs_len != 0) {
    output_refs = (amduat_reference_t *)calloc(outputs_len, sizeof(*output_refs));
    if (output_refs == NULL) {
      trace_ok = false;
      goto cleanup;
    }
    for (i = 0; i < outputs_len; ++i) {
      store_err = amduat_asl_store_put(store, outputs[i], &output_refs[i]);
      if (store_err != AMDUAT_ASL_STORE_OK) {
        trace_ok = false;
        goto cleanup;
      }
    }
  }

  trace_ok = false;
  if (exec_invoked) {
    trace_ok = amduat_store_trace(
        store, scheme_ref, program_ref, input_refs, input_refs_len,
        has_params_ref, params_ref, &core_result,
        program_decoded ? &program : NULL,
        program_decoded ? &trace_eval : NULL, &trace_ref);
    if (!trace_ok) {
      goto cleanup;
    }
  }

  if (!amduat_store_surface_result(
          store, &core_result, scheme_ref, program_ref, input_refs,
          input_refs_len, output_refs, output_refs_len, has_params_ref,
          params_ref, false, NULL, trace_ok, trace_ref, out_result_ref)) {
    goto cleanup;
  }

  *out_output_refs = output_refs;
  *out_output_refs_len = output_refs_len;

  if (program_decoded) {
    amduat_enc_pel_program_dag_free(&program);
    amduat_pel_program_dag_trace_free(&trace_eval);
  }
  for (i = 0; i < input_refs_len; ++i) {
    amduat_artifact_free(&input_artifacts[i]);
  }
  free(input_artifacts);
  if (has_params_artifact) {
    amduat_artifact_free(&params_artifact);
  }
  amduat_pel_program_dag_free_outputs(outputs, outputs_len);
  amduat_pel_execution_result_free(&core_result);
  return true;

cleanup:
  amduat_pel_surf_free_refs(output_refs, output_refs_len);
  if (program_decoded) {
    amduat_enc_pel_program_dag_free(&program);
    amduat_pel_program_dag_trace_free(&trace_eval);
  }
  for (i = 0; i < input_refs_len; ++i) {
    amduat_artifact_free(&input_artifacts[i]);
  }
  free(input_artifacts);
  if (has_params_artifact) {
    amduat_artifact_free(&params_artifact);
  }
  amduat_pel_program_dag_free_outputs(outputs, outputs_len);
  amduat_pel_execution_result_free(&core_result);
  return false;
}