#include "amduat/pel/opreg_kernel.h"
#include "amduat/pel/opreg_kernel_params.h"
#include "amduat/hash/asl1.h"

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

enum { AMDUAT_PEL_KERNEL_STATUS_INTERNAL = 1 };

static bool amduat_name_eq(amduat_octets_t name, const char *literal) {
  size_t len = strlen(literal);
  if (name.len != len) {
    return false;
  }
  if (len == 0) {
    return true;
  }
  return name.data != NULL &&
         memcmp(name.data, literal, len) == 0;
}

static const amduat_pel_kernel_op_desc_t k_concat_desc = {
    AMDUAT_PEL_KERNEL_OP_CONCAT,
    AMDUAT_PEL_KERNEL_OP_CODE_CONCAT,
    1,
    SIZE_MAX,
    1
};

static const amduat_pel_kernel_op_desc_t k_slice_desc = {
    AMDUAT_PEL_KERNEL_OP_SLICE,
    AMDUAT_PEL_KERNEL_OP_CODE_SLICE,
    1,
    1,
    1
};

static const amduat_pel_kernel_op_desc_t k_const_desc = {
    AMDUAT_PEL_KERNEL_OP_CONST,
    AMDUAT_PEL_KERNEL_OP_CODE_CONST,
    0,
    0,
    1
};

static const amduat_pel_kernel_op_desc_t k_hash_desc = {
    AMDUAT_PEL_KERNEL_OP_HASH_ASL1,
    AMDUAT_PEL_KERNEL_OP_CODE_HASH_ASL1,
    1,
    1,
    1
};

const amduat_pel_kernel_op_desc_t *amduat_pel_kernel_op_lookup(
    amduat_octets_t name,
    uint32_t version) {
  if (version != 1) {
    return NULL;
  }
  if (amduat_name_eq(name, AMDUAT_PEL_KERNEL_OP_CONCAT_NAME)) {
    return &k_concat_desc;
  }
  if (amduat_name_eq(name, AMDUAT_PEL_KERNEL_OP_SLICE_NAME)) {
    return &k_slice_desc;
  }
  if (amduat_name_eq(name, AMDUAT_PEL_KERNEL_OP_CONST_NAME)) {
    return &k_const_desc;
  }
  if (amduat_name_eq(name, AMDUAT_PEL_KERNEL_OP_HASH_ASL1_NAME)) {
    return &k_hash_desc;
  }
  return NULL;
}

static bool amduat_type_tags_match(const amduat_artifact_t *a,
                                   const amduat_artifact_t *b) {
  if (a->has_type_tag != b->has_type_tag) {
    return false;
  }
  if (!a->has_type_tag) {
    return true;
  }
  return a->type_tag.tag_id == b->type_tag.tag_id;
}

static bool amduat_alloc_outputs(size_t count,
                                 amduat_artifact_t **out_outputs) {
  amduat_artifact_t *outputs;

  if (out_outputs == NULL) {
    return false;
  }
  outputs = (amduat_artifact_t *)calloc(count, sizeof(*outputs));
  if (outputs == NULL) {
    return false;
  }
  *out_outputs = outputs;
  return true;
}

static bool amduat_kernel_concat(const amduat_artifact_t *inputs,
                                 size_t inputs_len,
                                 amduat_artifact_t **out_outputs,
                                 size_t *out_outputs_len,
                                 uint32_t *out_status_code) {
  size_t i;
  size_t total_len = 0;
  uint8_t *buffer;

  if (inputs == NULL) {
    if (out_status_code) {
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
    }
    return false;
  }
  if (inputs_len == 0) {
    if (out_status_code) {
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
    }
    return false;
  }

  for (i = 1; i < inputs_len; ++i) {
    if (!amduat_type_tags_match(&inputs[0], &inputs[i])) {
      if (out_status_code) {
        *out_status_code = AMDUAT_PEL_KERNEL_STATUS_CONCAT_TYPE_TAG_MISMATCH;
      }
      return false;
    }
  }

  for (i = 0; i < inputs_len; ++i) {
    if (inputs[i].bytes.len != 0 && inputs[i].bytes.data == NULL) {
      if (out_status_code) {
        *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
      }
      return false;
    }
    if (inputs[i].bytes.len > SIZE_MAX - total_len) {
      if (out_status_code) {
        *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
      }
      return false;
    }
    total_len += inputs[i].bytes.len;
  }

  buffer = NULL;
  if (total_len != 0) {
    buffer = (uint8_t *)malloc(total_len);
    if (buffer == NULL) {
      if (out_status_code) {
        *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
      }
      return false;
    }
  }

  {
    size_t offset = 0;
    for (i = 0; i < inputs_len; ++i) {
      if (inputs[i].bytes.len != 0 && inputs[i].bytes.data != NULL) {
        memcpy(buffer + offset, inputs[i].bytes.data, inputs[i].bytes.len);
      }
      offset += inputs[i].bytes.len;
    }
  }

  if (!amduat_alloc_outputs(1, out_outputs)) {
    free(buffer);
    if (out_status_code) {
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
    }
    return false;
  }

  (*out_outputs)[0].bytes = amduat_octets(buffer, total_len);
  if (inputs[0].has_type_tag) {
    (*out_outputs)[0].has_type_tag = true;
    (*out_outputs)[0].type_tag = inputs[0].type_tag;
  } else {
    (*out_outputs)[0].has_type_tag = false;
    (*out_outputs)[0].type_tag.tag_id = 0;
  }
  if (out_outputs_len) {
    *out_outputs_len = 1;
  }
  return true;
}

static bool amduat_kernel_slice(const amduat_artifact_t *inputs,
                                size_t inputs_len,
                                const amduat_pel_kernel_slice_params_t *params,
                                amduat_artifact_t **out_outputs,
                                size_t *out_outputs_len,
                                uint32_t *out_status_code) {
  const amduat_artifact_t *input;
  uint64_t offset;
  uint64_t length;
  uint64_t total_len;
  uint8_t *buffer;

  if (inputs == NULL) {
    if (out_status_code) {
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
    }
    return false;
  }
  if (inputs_len != 1 || params == NULL) {
    if (out_status_code) {
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
    }
    return false;
  }

  input = &inputs[0];
  total_len = (uint64_t)input->bytes.len;
  offset = params->offset;
  length = params->length;
  if (offset > total_len || length > total_len - offset) {
    if (out_status_code) {
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_SLICE_RANGE_OUT_OF_BOUNDS;
    }
    return false;
  }

  buffer = NULL;
  if (length != 0) {
    buffer = (uint8_t *)malloc((size_t)length);
    if (buffer == NULL) {
      if (out_status_code) {
        *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
      }
      return false;
    }
    if (input->bytes.data == NULL) {
      free(buffer);
      if (out_status_code) {
        *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
      }
      return false;
    }
    memcpy(buffer, input->bytes.data + (size_t)offset, (size_t)length);
  }

  if (!amduat_alloc_outputs(1, out_outputs)) {
    free(buffer);
    if (out_status_code) {
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
    }
    return false;
  }

  (*out_outputs)[0].bytes = amduat_octets(buffer, (size_t)length);
  if (input->has_type_tag) {
    (*out_outputs)[0].has_type_tag = true;
    (*out_outputs)[0].type_tag = input->type_tag;
  } else {
    (*out_outputs)[0].has_type_tag = false;
    (*out_outputs)[0].type_tag.tag_id = 0;
  }
  if (out_outputs_len) {
    *out_outputs_len = 1;
  }
  return true;
}

static bool amduat_kernel_const(const amduat_pel_kernel_const_params_t *params,
                                amduat_artifact_t **out_outputs,
                                size_t *out_outputs_len,
                                uint32_t *out_status_code) {
  uint8_t *buffer;

  if (params == NULL) {
    if (out_status_code) {
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
    }
    return false;
  }
  if (params->bytes.len != 0 && params->bytes.data == NULL) {
    if (out_status_code) {
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
    }
    return false;
  }

  buffer = NULL;
  if (params->bytes.len != 0) {
    buffer = (uint8_t *)malloc(params->bytes.len);
    if (buffer == NULL) {
      if (out_status_code) {
        *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
      }
      return false;
    }
    if (params->bytes.data != NULL) {
      memcpy(buffer, params->bytes.data, params->bytes.len);
    }
  }

  if (!amduat_alloc_outputs(1, out_outputs)) {
    free(buffer);
    if (out_status_code) {
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
    }
    return false;
  }

  (*out_outputs)[0].bytes = amduat_octets(buffer, params->bytes.len);
  if (params->has_type_tag) {
    (*out_outputs)[0].has_type_tag = true;
    (*out_outputs)[0].type_tag.tag_id = params->tag_id;
  } else {
    (*out_outputs)[0].has_type_tag = false;
    (*out_outputs)[0].type_tag.tag_id = 0;
  }
  if (out_outputs_len) {
    *out_outputs_len = 1;
  }
  return true;
}

static bool amduat_kernel_hash(const amduat_artifact_t *inputs,
                               size_t inputs_len,
                               const amduat_pel_kernel_hash_params_t *params,
                               amduat_artifact_t **out_outputs,
                               size_t *out_outputs_len,
                               uint32_t *out_status_code) {
  const amduat_hash_asl1_desc_t *desc;
  uint8_t *buffer;

  if (inputs == NULL) {
    if (out_status_code) {
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
    }
    return false;
  }
  if (inputs_len != 1 || params == NULL) {
    if (out_status_code) {
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
    }
    return false;
  }

  desc = amduat_hash_asl1_desc_lookup(params->hash_id);
  if (desc == NULL || desc->digest_len == 0) {
    if (out_status_code) {
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
    }
    return false;
  }
  if (inputs[0].bytes.len != 0 && inputs[0].bytes.data == NULL) {
    if (out_status_code) {
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
    }
    return false;
  }

  buffer = (uint8_t *)malloc(desc->digest_len);
  if (buffer == NULL) {
    if (out_status_code) {
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
    }
    return false;
  }

  if (!amduat_hash_asl1_digest(params->hash_id, inputs[0].bytes, buffer,
                               desc->digest_len)) {
    free(buffer);
    if (out_status_code) {
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
    }
    return false;
  }

  if (!amduat_alloc_outputs(1, out_outputs)) {
    free(buffer);
    if (out_status_code) {
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
    }
    return false;
  }

  (*out_outputs)[0].bytes = amduat_octets(buffer, desc->digest_len);
  (*out_outputs)[0].has_type_tag = false;
  (*out_outputs)[0].type_tag.tag_id = 0;
  if (out_outputs_len) {
    *out_outputs_len = 1;
  }
  return true;
}

bool amduat_pel_kernel_op_eval(
    const amduat_pel_kernel_op_desc_t *desc,
    const amduat_artifact_t *inputs,
    size_t inputs_len,
    const amduat_pel_kernel_params_t *params,
    amduat_artifact_t **out_outputs,
    size_t *out_outputs_len,
    uint32_t *out_status_code) {
  if (out_outputs_len != NULL) {
    *out_outputs_len = 0;
  }
  if (out_outputs != NULL) {
    *out_outputs = NULL;
  }

  if (desc == NULL || out_status_code == NULL) {
    return false;
  }
  if (params == NULL) {
    *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
    return false;
  }

  switch (desc->kind) {
    case AMDUAT_PEL_KERNEL_OP_CONCAT:
      return amduat_kernel_concat(inputs, inputs_len, out_outputs,
                                  out_outputs_len, out_status_code);
    case AMDUAT_PEL_KERNEL_OP_SLICE:
      return amduat_kernel_slice(inputs, inputs_len, &params->value.slice,
                                 out_outputs, out_outputs_len, out_status_code);
    case AMDUAT_PEL_KERNEL_OP_CONST:
      return amduat_kernel_const(&params->value.konst, out_outputs,
                                 out_outputs_len, out_status_code);
    case AMDUAT_PEL_KERNEL_OP_HASH_ASL1:
      return amduat_kernel_hash(inputs, inputs_len, &params->value.hash,
                                out_outputs, out_outputs_len, out_status_code);
    default:
      *out_status_code = AMDUAT_PEL_KERNEL_STATUS_INTERNAL;
      return false;
  }
}