#include "sha2.h"
#include "../trace.h"
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>

#ifdef TEST
// if you want to get the internals of sha2 printed, uncomment the line below
// #define SHOW_INTERNALS
#endif // TEST

#ifdef SHOW_INTERNALS
#define sha2_256_show_internal_state(context)                                  \
  TRACE("Internal State:\n\tcorrupted: %d\n\tintermediate_hash: ",             \
        context->corrupted);                                                   \
  dump_hash(context);                                                          \
  TRACE("\tmessage_block_index: %d\n\tmessage_block: ",                        \
        context->message_block_index);                                         \
  dump_data((uint8_t *)context->message_block, SHA2_256_BlockSize);            \
  TRACE("\tcorrupted: %d\n\tcomputed: %d\n", context->corrupted,               \
        context->computed);
#else
#define sha2_256_show_internal_state(context)
#endif // SHOW_INTERNALS

#define err_handler(RES)                                                       \
  if (RES) {                                                                   \
    return RES;                                                                \
  }

static uint32_t addTemp = 0;
#define sha2_256_add_length(context, length)                                   \
  (addTemp = (context)->length_low,                                            \
   (context)->corrupted = (((context)->length_low += (length)) < addTemp) &&   \
                                  (++(context)->length_high == 0)              \
                              ? shaInputTooLong                                \
                              : (context)->corrupted)

#define dump_hash(context)                                                     \
  for (int o = 0; o < SHA2_256_HashParts; o++) {                               \
    TRACE("%08x", context->intermediate_hash[o]);                              \
  }                                                                            \
  TRACELN("")

/* Define the SHA shift, rotate left, and rotate right macros */
#define SHA256_SHR(bits, word) ((word) >> (bits))
#define SHA256_ROTL(bits, word) (((word) << (bits)) | ((word) >> (32 - (bits))))
#define SHA256_ROTR(bits, word) (((word) >> (bits)) | ((word) << (32 - (bits))))

/* Define the SHA SIGMA and sigma macros, yes these have the same name and are
 * different */
#define SHA256_SIGMA0(word)                                                    \
  (SHA256_ROTR(2, word) ^ SHA256_ROTR(13, word) ^ SHA256_ROTR(22, word))
#define SHA256_SIGMA1(word)                                                    \
  (SHA256_ROTR(6, word) ^ SHA256_ROTR(11, word) ^ SHA256_ROTR(25, word))
#define SHA256_sigma0(word)                                                    \
  (SHA256_ROTR(7, word) ^ SHA256_ROTR(18, word) ^ SHA256_SHR(3, word))
#define SHA256_sigma1(word)                                                    \
  (SHA256_ROTR(17, word) ^ SHA256_ROTR(19, word) ^ SHA256_SHR(10, word))

#define SHA_Ch(x, y, z) (((x) & ((y) ^ (z))) ^ (z))
#define SHA_Maj(x, y, z) (((x) & ((y) | (z))) | ((y) & (z)))

static const uint32_t H[SHA2_256_HashParts] = {
    0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a,
    0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19};
static const uint32_t K[SHA2_256_BlockSize] = {
    0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1,
    0x923f82a4, 0xab1c5ed5, 0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3,
    0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174, 0xe49b69c1, 0xefbe4786,
    0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
    0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147,
    0x06ca6351, 0x14292967, 0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13,
    0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85, 0xa2bfe8a1, 0xa81a664b,
    0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
    0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a,
    0x5b9cca4f, 0x682e6ff3, 0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208,
    0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2};

/*
 * Description:
 *   This helper function will process the next 512 bits of the
 *   message stored in the message_block array.
 *
 * Parameters:
 *   context: [in/out]
 *     The SHA context to update.
 *
 * Returns:
 *   Nothing.
 *
 * Comments:
 *   Many of the variable names in this code, especially the
 *   single character names, were used because those were the
 *   names used in the Secure Hash Standard.
 */
static void sha2_256_process_message_block(SHA2Context *context) {
#define dump_letter(X) TRACE("%s: %08x ", STR(X), X)
#define dump_alphabet                                                          \
  dump_letter(A) dump_letter(B) dump_letter(C) dump_letter(D) dump_letter(E)   \
      dump_letter(F) dump_letter(G) dump_letter(H) TRACELN("")
#define dump_w                                                                 \
  TRACELN("Dumping W:");                                                       \
  for (int o = 0; o < SHA2_256_BlockSize; o++) {                               \
    TRACE("\tW[%d]: %032b\n", o, W[o])                                         \
  }

#ifdef SHOW_INTERNALS
  TRACELN("before process message block");
  sha2_256_show_internal_state(context);
#endif // SHOW_INTERNALS

  uint8_t t, t4;                   /* Loop counter */
  uint32_t temp1, temp2;           /* Temporary word value */
  uint32_t W[SHA2_256_BlockSize];  /* Word sequence */
  uint32_t A, B, C, D, E, F, G, H; /* Word buffers */

  /*
   * Initialize the first 16 words in the array W
   */
  for (t = t4 = 0; t < 16; t++, t4 += 4) {
    W[t] = (((uint32_t)context->message_block[t4]) << 24) |
           (((uint32_t)context->message_block[t4 + 1]) << 16) |
           (((uint32_t)context->message_block[t4 + 2]) << 8) |
           (((uint32_t)context->message_block[t4 + 3]));
  }

#ifdef SHOW_INTERNALS
  dump_w;
#endif // SHOW_INTERNALS

  for (t = 16; t < SHA2_256_BlockSize; t++) {
    W[t] = SHA256_sigma1(W[t - 2]) + W[t - 7] + SHA256_sigma0(W[t - 15]) +
           W[t - 16];
  }

#ifdef SHOW_INTERNALS
  dump_w;
#endif // SHOW_INTERNALS

  A = context->intermediate_hash[0];
  B = context->intermediate_hash[1];
  C = context->intermediate_hash[2];
  D = context->intermediate_hash[3];
  E = context->intermediate_hash[4];
  F = context->intermediate_hash[5];
  G = context->intermediate_hash[6];
  H = context->intermediate_hash[7];

  for (t = 0; t < SHA2_256_BlockSize; t++) {
#ifdef SHOW_INTERNALS
    TRACE("Iter: %d\t", t);
    dump_alphabet;
#endif // SHOW_INTERNALS

    temp1 = H + SHA256_SIGMA1(E) + SHA_Ch(E, F, G) + K[t] + W[t];
    temp2 = SHA256_SIGMA0(A) + SHA_Maj(A, B, C);
    H = G;
    G = F;
    F = E;
    E = D + temp1;
    D = C;
    C = B;
    B = A;
    A = temp1 + temp2;
  }

#ifdef SHOW_INTERNALS
  TRACE("DONE:\t")
  dump_alphabet;

  TRACE("Intermediate hash before we add the working variables:\n");
  for (int o = 0; o < SHA2_256_HashParts; o++) {
    TRACE("\t%d:\t%032b\n", o, context->intermediate_hash[o]);
  }
#endif // SHOW_INTERNALS

  context->intermediate_hash[0] += A;
  context->intermediate_hash[1] += B;
  context->intermediate_hash[2] += C;
  context->intermediate_hash[3] += D;
  context->intermediate_hash[4] += E;
  context->intermediate_hash[5] += F;
  context->intermediate_hash[6] += G;
  context->intermediate_hash[7] += H;

  context->message_block_index = 0;

#ifdef SHOW_INTERNALS
  TRACE("Hash data if we finish now: ");
  dump_hash(context);
  TRACELN("\nafter process message block");
  sha2_256_show_internal_state(context);
#endif // SHOW_INTERNALS
}

/*
 * Description:
 *   According to the standard, the message must be padded to the next
 *   even multiple of 512 bits.  The first padding bit must be a '1'.
 *   The last 64 bits represent the length of the original message.
 *   All bits in between should be 0.  This helper function will pad
 *   the message according to those rules by filling the
 *   Message_Block array accordingly.  When it returns, it can be
 *   assumed that the message digest has been computed.
 *
 * Parameters:
 *   context: [in/out]
 *     The context to pad.
 *   Pad_Byte: [in]
 *     The last byte to add to the message block before the 0-padding
 *     and length.  This will contain the last bits of the message
 *     followed by another single bit.  If the message was an
 *     exact multiple of 8-bits long, Pad_Byte will be 0x80.
 *
 * Returns:
 *   Nothing.
 */
static void sha2_256_pad_message(SHA2Context *context, uint8_t pad_byte) {
  /*
   * Check to see if the current message block is too small to hold
   * the initial padding bits and length.  If so, we will pad the
   * block, process it, and then continue padding into a second
   * block.
   */
  if (context->message_block_index >= (SHA2_256_BlockSize - 8)) {
    context->message_block[context->message_block_index++] = pad_byte;
    while (context->message_block_index < SHA2_256_BlockSize)
      context->message_block[context->message_block_index++] = 0;
    sha2_256_process_message_block(context);
  } else
    context->message_block[context->message_block_index++] = pad_byte;

  while (context->message_block_index < (SHA2_256_BlockSize - 8))
    context->message_block[context->message_block_index++] = 0;

  /*
   * store the message length as the last 8 octets
   */
  context->message_block[56] = (uint8_t)(context->length_high >> 24);
  context->message_block[57] = (uint8_t)(context->length_high >> 16);
  context->message_block[58] = (uint8_t)(context->length_high >> 8);
  context->message_block[59] = (uint8_t)(context->length_high);
  context->message_block[60] = (uint8_t)(context->length_low >> 24);
  context->message_block[61] = (uint8_t)(context->length_low >> 16);
  context->message_block[62] = (uint8_t)(context->length_low >> 8);
  context->message_block[63] = (uint8_t)(context->length_low);

  sha2_256_process_message_block(context);

  sha2_256_show_internal_state(context);
}

static void sha2_256_finalize(SHA2Context *context, uint8_t pad_byte) {
  sha2_256_pad_message(context, pad_byte);
}

/*
 *  sha2_reset
 *
 *  Description:
 *      This function will initialize the SHA2Context in preparation
 *      for computing a new SHA2 message digest.
 *
 *  Parameters:
 *      context: [in/out]
 *          The context to reset.
 *
 *  Returns:
 *      SHA2 Error Code.
 *
 */
SHA2Result sha2_256_reset(SHA2Context *context) {
  size_t i;

  if (!context) {
    return shaNull;
  }

  context->length_low = 0;
  context->length_high = 0;
  context->message_block_index = 0;

  for (i = 0; i < SHA2_256_BlockSize; i++) {
    context->message_block[i] = 0;
  }
  for (i = 0; i < SHA2_256_HashParts; i++) {
    context->intermediate_hash[i] = H[i];
  }

  context->corrupted = 0;
  context->computed = 0;

  return shaSuccess;
}

/*
 *  sha2_input
 *
 *  Description:
 *      This function accepts an array of bytes as the next portion
 *      of the message.
 *
 *  Parameters:
 *      context: [in/out]
 *          The SHA context to update
 *      data: [in]
 *          An array of characters representing the next portion of
 *          the message.
 *      length: [in]
 *          The length of the message in message_array
 *
 *  Returns:
 *      sha Error Code.
 *
 */
SHA2Result sha2_256_input(SHA2Context *context, const uint8_t data[],
                          size_t len) {
  sha2_256_show_internal_state(context);

  if (!context) {
    TRACELN("Context was NULL");
    return shaNull;
  }
  if (!len)
    return shaSuccess;
  if (!data) {
    TRACELN("Data was NULL");
    return shaNull;
  }
  if (context->computed) {
    TRACELN("context was already computed");
    context->corrupted = shaStateError;
  }
  if (context->corrupted) {
    TRACELN("context is corrupted");
    return context->corrupted;
  }

  while (len--) {
    context->message_block[context->message_block_index++] = *data;

    if ((sha2_256_add_length(context, 8) == shaSuccess) &&
        (context->message_block_index == SHA2_256_BlockSize)) {
      TRACELN("processing the message block after inputting")
      sha2_256_process_message_block(context);
    }

    data++;
  }

  return context->corrupted;
}

/*
 *  sha2_result
 *
 * Description:
 *   This function will return the 256-bit message digest
 *   into the digest array provided by the caller.
 *   NOTE:
 *    The first octet of hash is stored in the element with index 0,
 *    the last octet of hash in the element with index 31.
 *
 * Parameters:
 *   context: [in/out]
 *     The context to use to calculate the SHA hash.
 *   digest: [out]
 *     Where the digest is returned.
 *
 * Returns:
 *   sha Error Code.
 *
 */
SHA2Result sha2_256_result(SHA2Context *context, SHA2Digest digest) {
  size_t i;

  if (!context)
    return shaNull;
  if (!digest)
    return shaNull;
  if (context->corrupted)
    return context->corrupted;

  if (!context->computed) {
    sha2_256_finalize(context, 0x80);
  }

  for (i = 0; i < SHA2_256_HashParts; i++)
    digest[i] = context->intermediate_hash[i];

  return shaSuccess;
}

SHA2Result sha2_256_oneshot(const uint8_t data[], const size_t len,
                            SHA2Digest digest) {
  SHA2Context context;
  SHA2Result res;

  res = sha2_256_reset(&context);
  err_handler(res);
  sha2_256_show_internal_state((&context));

  res = sha2_256_input(&context, data, len);
  err_handler(res);
  sha2_256_show_internal_state((&context));

  res = sha2_256_result(&context, digest);
  err_handler(res);
  sha2_256_show_internal_state((&context));

  return res;
}

void sha2_256_oneshot_n_times(uint32_t n, const uint8_t data[],
                              const size_t len) {
  SHA2Digest digest;
  for (uint32_t i = 0; i < n; i++) {
    sha2_256_oneshot(data, len, digest);
  }
}
