#include "../../lib/md2_impls/md2_1.h"

#include "../../lib/md2_impls/md2_common.h"
#include <immintrin.h>


void md2_checksum_1(size_t len, uint8_t *buf)
{
  uint8_t l = 0;

  for (size_t i = 0; i < len / 16; i++)
  {
    for (int j = 0; j < 16; j++)
    {
      u_int8_t c = buf[i * 16 + j];
      // reference is wrong. It says: Set C[j] to S[c xor L]. But it should be:
      buf[len + j] ^= MD2_PI_SUBST[c ^ l];
      l = buf[len + j];
    }
  }
}

static uint8_t PADDING[17][16] = {
    {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
    {1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
    {2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
    {3, 3, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
    {4, 4, 4, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
    {5, 5, 5, 5, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
    {6, 6, 6, 6, 6, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
    {7, 7, 7, 7, 7, 7, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0},
    {8, 8, 8, 8, 8, 8, 8, 8, 0, 0, 0, 0, 0, 0, 0, 0},
    {9, 9, 9, 9, 9, 9, 9, 9, 9, 0, 0, 0, 0, 0, 0, 0},
    {10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 0, 0, 0, 0, 0, 0},
    {11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 0, 0, 0, 0, 0},
    {12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 0, 0, 0, 0},
    {13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 0, 0, 0},
    {14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 0, 0},
    {15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 0},
    {16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16}};

void md2_hash_1(size_t len, const uint8_t buf[len], uint8_t out[16])
{
  // === step 1 ===
  int paddingNeeded = 16 - (len & 15);
  len += paddingNeeded;

  uint8_t* newBuf = aligned_alloc(16, sizeof(uint8_t)*(len + 16));
  if (newBuf == NULL) {
    printf("Failure: aligned_alloc returns NULL");
    return -1;
  }

  for(size_t i = 0; i < 16; i++) {
    newBuf[len + i] = 0;
  }

  memcpy(newBuf, buf, len - paddingNeeded);

  memcpy(newBuf + len - paddingNeeded, PADDING + paddingNeeded, paddingNeeded);

  // === step 2 ===
  md2_checksum_1(len, newBuf);

  // === step 3 ===
  uint8_t *messageDigestBuf = aligned_alloc(16, sizeof(uint8_t) * 48);
  if (messageDigestBuf == NULL) {
    printf("Failure: aligned_alloc returns NULL");
    return -1;
  }

  for (size_t i = 0; i < 48; i++) {
    messageDigestBuf[i] = 0;
  }

  // === step 4 ===
  __m128i vx;
  __m128i vy;
  for (size_t i = 0; i <= (len + 16) / 16 - 1; i++)
  {
    vx = _mm_load_si128((__m128i*) (newBuf + i * 16));
    _mm_store_si128((__m128i*) (messageDigestBuf + 16), vx);
    vy = _mm_load_si128((__m128i*) (messageDigestBuf));
    vy = _mm_xor_si128(vy, vx);
    _mm_store_si128((__m128i*) (messageDigestBuf + 32), vy);
  

    u_int8_t t = 0;

    for (int j = 0; j < 18; j++)
    {
      for (int k = 0; k < 48; k++)
      {
        t = messageDigestBuf[k] = messageDigestBuf[k] ^ MD2_PI_SUBST[t];
      }
      t = (t + j) & 255;
    }
  }

  memcpy(out, messageDigestBuf, 16);

  free(messageDigestBuf);
  free(newBuf);
}