#include "amduat/pel/derivation_sid.h"

#include "amduat/enc/asl1_core_codec.h"
#include "amduat/hash/asl1.h"

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

static bool amduat_derivation_sid_check_ref(amduat_reference_t ref) {
  const amduat_hash_asl1_desc_t *desc;
  size_t digest_len;

  if (ref.digest.len != 0 && ref.digest.data == NULL) {
    return false;
  }
  if (amduat_hash_asl1_is_reserved(ref.hash_id)) {
    return false;
  }

  desc = amduat_hash_asl1_desc_lookup(ref.hash_id);
  digest_len = ref.digest.len;
  if (desc != NULL && desc->digest_len != 0) {
    assert(desc->digest_len == digest_len);
    if (desc->digest_len != digest_len) {
      return false;
    }
  }

  return true;
}

bool amduat_pel_derivation_sid_compute(
    const amduat_pel_derivation_sid_input_t *in,
    amduat_octets_t *out_sid) {
  amduat_hash_asl1_stream_t stream;
  uint8_t marker;
  size_t i;

  if (out_sid == NULL) {
    return false;
  }
  out_sid->data = NULL;
  out_sid->len = 0;

  if (in == NULL) {
    return false;
  }
  if (in->input_refs_len != 0 && in->input_refs == NULL) {
    return false;
  }
  if (in->has_params_ref && !amduat_derivation_sid_check_ref(in->params_ref)) {
    return false;
  }
  if (!amduat_derivation_sid_check_ref(in->program_ref)) {
    return false;
  }
  for (i = 0; i < in->input_refs_len; ++i) {
    if (!amduat_derivation_sid_check_ref(in->input_refs[i])) {
      return false;
    }
  }

  if (!amduat_hash_asl1_stream_init(AMDUAT_HASH_ASL1_ID_SHA256, &stream)) {
    return false;
  }

  {
    amduat_octets_t ref_bytes = amduat_octets(NULL, 0u);
    bool ok;

    /* ReferenceBytes are self-delimiting because hash_id implies digest_len. */
    ok = amduat_enc_asl1_core_encode_reference_v1(in->program_ref, &ref_bytes);
    if (!ok ||
        !amduat_hash_asl1_stream_update(&stream, ref_bytes.data,
                                        ref_bytes.len)) {
      amduat_hash_asl1_stream_destroy(&stream);
      free((void *)ref_bytes.data);
      return false;
    }
    free((void *)ref_bytes.data);
  }

  for (i = 0; i < in->input_refs_len; ++i) {
    amduat_octets_t ref_bytes = amduat_octets(NULL, 0u);
    bool ok = amduat_enc_asl1_core_encode_reference_v1(in->input_refs[i],
                                                       &ref_bytes);
    if (!ok ||
        !amduat_hash_asl1_stream_update(&stream, ref_bytes.data,
                                        ref_bytes.len)) {
      amduat_hash_asl1_stream_destroy(&stream);
      free((void *)ref_bytes.data);
      return false;
    }
    free((void *)ref_bytes.data);
  }

  marker = in->has_params_ref ? 0x01u : 0x00u;
  if (!amduat_hash_asl1_stream_update(&stream, &marker, 1u)) {
    amduat_hash_asl1_stream_destroy(&stream);
    return false;
  }
  if (in->has_params_ref) {
    amduat_octets_t ref_bytes = amduat_octets(NULL, 0u);
    bool ok = amduat_enc_asl1_core_encode_reference_v1(in->params_ref,
                                                       &ref_bytes);
    if (!ok ||
        !amduat_hash_asl1_stream_update(&stream, ref_bytes.data,
                                        ref_bytes.len)) {
      amduat_hash_asl1_stream_destroy(&stream);
      free((void *)ref_bytes.data);
      return false;
    }
    free((void *)ref_bytes.data);
  }

  /* ExecProfile is currently absent; presence must be explicitly marked. */
  marker = in->has_exec_profile ? 0x01u : 0x00u;
  if (!amduat_hash_asl1_stream_update(&stream, &marker, 1u)) {
    amduat_hash_asl1_stream_destroy(&stream);
    return false;
  }
  if (in->has_exec_profile) {
    if (in->exec_profile.len != 0 && in->exec_profile.data == NULL) {
      amduat_hash_asl1_stream_destroy(&stream);
      return false;
    }
    if (!amduat_hash_asl1_stream_update(&stream, in->exec_profile.data,
                                        in->exec_profile.len)) {
      amduat_hash_asl1_stream_destroy(&stream);
      return false;
    }
  }

  {
    const amduat_hash_asl1_desc_t *desc =
        amduat_hash_asl1_desc_lookup(AMDUAT_HASH_ASL1_ID_SHA256);
    uint8_t *digest;
    size_t digest_len;

    if (desc == NULL || desc->digest_len == 0u) {
      amduat_hash_asl1_stream_destroy(&stream);
      return false;
    }
    digest_len = desc->digest_len;
    digest = (uint8_t *)malloc(digest_len);
    if (digest == NULL) {
      amduat_hash_asl1_stream_destroy(&stream);
      return false;
    }
    if (!amduat_hash_asl1_stream_final(&stream, digest, digest_len)) {
      free(digest);
      amduat_hash_asl1_stream_destroy(&stream);
      return false;
    }
    amduat_hash_asl1_stream_destroy(&stream);
    *out_sid = amduat_octets(digest, digest_len);
  }

  return true;
}