#include "amduat/hash/asl1.h"

#include "sha256.h"

#include <stdlib.h>
#include <string.h>

static amduat_hash_asl1_desc_t g_hash_asl1_descs[] = {
    {AMDUAT_HASH_ASL1_ID_SHA256,
     "HASH-ASL1-256",
     32,
     {amduat_hash_asl1_sha256_digest, NULL}},
    {0x0002, "HASH-ASL1-512", 64, {NULL, NULL}},
    {0x8001, "HASH-ASL1-PQ1", 0, {NULL, NULL}}};

typedef struct {
  amduat_hash_id_t start;
  amduat_hash_id_t end;
} amduat_hash_asl1_reserved_range_t;

static const amduat_hash_asl1_reserved_range_t
    g_hash_asl1_reserved_ranges[] = {
        {0x0000, 0x0000},
        {0x0002, 0x0002},
        {0x8001, 0x8001},
        {0x8002, 0x80FF},
    };

#ifdef AMDUAT_HASH_ASL1_ENABLE_SHA256_STUB
static bool amduat_hash_asl1_sha256_stub(void *ctx,
                                         amduat_octets_t input,
                                         uint8_t *out,
                                         size_t out_len) {
  size_t i;

  (void)ctx;
  (void)input;

  if (out == NULL || out_len < 32) {
    return false;
  }

  for (i = 0; i < 32; ++i) {
    out[i] = 0;
  }

  return true;
}
#endif

amduat_octets_t amduat_hash_asl1_key(amduat_hash_id_t hash_id, uint8_t out[2]) {
  out[0] = (uint8_t)(hash_id >> 8);
  out[1] = (uint8_t)(hash_id & 0xff);
  return amduat_octets(out, 2);
}

const amduat_hash_asl1_desc_t *amduat_hash_asl1_desc_lookup(
    amduat_hash_id_t hash_id) {
  size_t i;

  if (amduat_hash_asl1_is_reserved(hash_id)) {
    return NULL;
  }

  for (i = 0; i < sizeof(g_hash_asl1_descs) / sizeof(g_hash_asl1_descs[0]);
       ++i) {
    if (g_hash_asl1_descs[i].hash_id == hash_id) {
      return &g_hash_asl1_descs[i];
    }
  }

  return NULL;
}

const amduat_hash_asl1_desc_t *amduat_hash_asl1_descs(size_t *out_len) {
  if (out_len != NULL) {
    *out_len = sizeof(g_hash_asl1_descs) / sizeof(g_hash_asl1_descs[0]);
  }
  return g_hash_asl1_descs;
}

bool amduat_hash_asl1_register_impl(amduat_hash_id_t hash_id,
                                    amduat_hash_asl1_impl_t impl) {
  size_t i;

  if (amduat_hash_asl1_is_reserved(hash_id)) {
    return false;
  }
  if (hash_id == AMDUAT_HASH_ASL1_ID_SHA256) {
    return false;
  }

  for (i = 0; i < sizeof(g_hash_asl1_descs) / sizeof(g_hash_asl1_descs[0]);
       ++i) {
    if (g_hash_asl1_descs[i].hash_id == hash_id) {
      g_hash_asl1_descs[i].impl = impl;
      return true;
    }
  }

  return false;
}

bool amduat_hash_asl1_digest(amduat_hash_id_t hash_id,
                             amduat_octets_t input,
                             uint8_t *out,
                             size_t out_len) {
  const amduat_hash_asl1_desc_t *desc;

  if (out == NULL) {
    return false;
  }

  desc = amduat_hash_asl1_desc_lookup(hash_id);
  if (desc == NULL || desc->digest_len == 0) {
    return false;
  }
  if (out_len < desc->digest_len) {
    return false;
  }
  if (desc->impl.digest == NULL) {
    return false;
  }

  return desc->impl.digest(desc->impl.ctx, input, out, desc->digest_len);
}

static bool amduat_hash_asl1_sha256_stream_update(void *ctx,
                                                  amduat_octets_t input) {
  if (input.len != 0 && input.data == NULL) {
    return false;
  }
  amduat_sha256_update((amduat_sha256_ctx_t *)ctx, input.data, input.len);
  return true;
}

static bool amduat_hash_asl1_sha256_stream_final(void *ctx,
                                                 uint8_t *out,
                                                 size_t out_len) {
  if (out == NULL || out_len < 32u) {
    return false;
  }
  amduat_sha256_final((amduat_sha256_ctx_t *)ctx, out);
  return true;
}

static void amduat_hash_asl1_sha256_stream_destroy(void *ctx) {
  free(ctx);
}

bool amduat_hash_asl1_stream_init(amduat_hash_id_t hash_id,
                                  amduat_hash_asl1_stream_t *out) {
  amduat_sha256_ctx_t *ctx;

  if (out == NULL) {
    return false;
  }
  memset(out, 0, sizeof(*out));
  if (amduat_hash_asl1_is_reserved(hash_id)) {
    return false;
  }
  if (hash_id != AMDUAT_HASH_ASL1_ID_SHA256) {
    return false;
  }

  ctx = (amduat_sha256_ctx_t *)malloc(sizeof(*ctx));
  if (ctx == NULL) {
    return false;
  }
  amduat_sha256_init(ctx);
  out->hash_id = hash_id;
  out->ctx = ctx;
  out->update = amduat_hash_asl1_sha256_stream_update;
  out->finalize = amduat_hash_asl1_sha256_stream_final;
  out->destroy = amduat_hash_asl1_sha256_stream_destroy;
  return true;
}

bool amduat_hash_asl1_stream_update(amduat_hash_asl1_stream_t *stream,
                                    amduat_octets_t input) {
  if (stream == NULL || stream->update == NULL) {
    return false;
  }
  return stream->update(stream->ctx, input);
}

bool amduat_hash_asl1_stream_final(amduat_hash_asl1_stream_t *stream,
                                   uint8_t *out,
                                   size_t out_len) {
  if (stream == NULL || stream->finalize == NULL) {
    return false;
  }
  return stream->finalize(stream->ctx, out, out_len);
}

void amduat_hash_asl1_stream_destroy(amduat_hash_asl1_stream_t *stream) {
  if (stream == NULL) {
    return;
  }
  if (stream->destroy != NULL) {
    stream->destroy(stream->ctx);
  }
  memset(stream, 0, sizeof(*stream));
}

bool amduat_hash_asl1_register_sha256_stub(void) {
#ifdef AMDUAT_HASH_ASL1_ENABLE_SHA256_STUB
  amduat_hash_asl1_impl_t impl;

  impl.digest = amduat_hash_asl1_sha256_stub;
  impl.ctx = NULL;
  return amduat_hash_asl1_register_impl(AMDUAT_HASH_ASL1_ID_SHA256, impl);
#else
  return false;
#endif
}

bool amduat_hash_asl1_is_reserved(amduat_hash_id_t hash_id) {
  size_t i;

  for (i = 0; i < sizeof(g_hash_asl1_reserved_ranges) /
                      sizeof(g_hash_asl1_reserved_ranges[0]);
       ++i) {
    if (hash_id >= g_hash_asl1_reserved_ranges[i].start &&
        hash_id <= g_hash_asl1_reserved_ranges[i].end) {
      return true;
    }
  }
  return false;
}