#include <stdint.h>
#include <string.h>
#include <stdio.h>
#include <inttypes.h>

#include "blake2/blake2.h"
#include "blake2/blake2-impl.h"

static const uint64_t blake2b_IV[8] = {
    UINT64_C(0x6a09e667f3bcc908), UINT64_C(0xbb67ae8584caa73b),
    UINT64_C(0x3c6ef372fe94f82b), UINT64_C(0xa54ff53a5f1d36f1),
    UINT64_C(0x510e527fade682d1), UINT64_C(0x9b05688c2b3e6c1f),
    UINT64_C(0x1f83d9abfb41bd6b), UINT64_C(0x5be0cd19137e2179)
};

static const unsigned int blake2b_sigma[12][16] = {
    {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15},
    {14, 10, 4, 8, 9, 15, 13, 6, 1, 12, 0, 2, 11, 7, 5, 3},
    {11, 8, 12, 0, 5, 2, 15, 13, 10, 14, 3, 6, 7, 1, 9, 4},
    {7, 9, 3, 1, 13, 12, 11, 14, 2, 6, 5, 10, 4, 0, 15, 8},
    {9, 0, 5, 7, 2, 4, 10, 15, 14, 1, 11, 12, 6, 8, 3, 13},
    {2, 12, 6, 10, 0, 11, 8, 3, 4, 13, 7, 5, 15, 14, 1, 9},
    {12, 5, 1, 15, 14, 13, 4, 10, 0, 7, 6, 3, 9, 2, 8, 11},
    {13, 11, 7, 14, 12, 1, 3, 9, 5, 0, 15, 4, 8, 6, 2, 10},
    {6, 15, 14, 9, 11, 3, 0, 8, 12, 2, 13, 7, 1, 4, 10, 5},
    {10, 2, 8, 4, 7, 6, 1, 5, 15, 11, 9, 14, 3, 12, 13, 0},
    {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15},
    {14, 10, 4, 8, 9, 15, 13, 6, 1, 12, 0, 2, 11, 7, 5, 3},
};

static BLAKE2_INLINE void blake2b_set_lastnode(blake2b_state *S) {
    S->f[1] = (uint64_t)-1;
}

static BLAKE2_INLINE void blake2b_set_lastblock(blake2b_state *S) {
    if (S->last_node) {
        blake2b_set_lastnode(S);
    }
    S->f[0] = (uint64_t)-1;
}

static BLAKE2_INLINE void blake2b_increment_counter(blake2b_state *S,
                                                    uint64_t inc) {
    S->t[0] += inc;
    S->t[1] += (S->t[0] < inc);
}

static BLAKE2_INLINE void blake2b_invalidate_state(blake2b_state *S) {
    burn(S, sizeof(*S));      /* wipe */
    blake2b_set_lastblock(S); /* invalidate for further use */
}

static BLAKE2_INLINE void blake2b_init0(blake2b_state *S) {
    memset(S, 0, sizeof(*S));
    memcpy(S->h, blake2b_IV, sizeof(S->h));
}


/*void print_state(blake2b_state BlakeHash) {
  printf(".h = {UINT64_C(%" PRIu64 "), UINT64_C(%" PRIu64 "),\n"
               "UINT64_C(%" PRIu64 "), UINT64_C(%" PRIu64 "),\n"
               "UINT64_C(%" PRIu64 "), UINT64_C(%" PRIu64 "),\n"
               "UINT64_C(%" PRIu64 "), UINT64_C(%" PRIu64 ")},\n"
         ".t = {UINT64_C(%" PRIu64 "), UINT64_C(%" PRIu64 ")},\n"
         ".f = {UINT64_C(%" PRIu64 "), UINT64_C(%" PRIu64 ")}\n",
         BlakeHash.h[0], BlakeHash.h[1], BlakeHash.h[2], BlakeHash.h[3],
         BlakeHash.h[4], BlakeHash.h[5], BlakeHash.h[6], BlakeHash.h[7],
         BlakeHash.t[0], BlakeHash.t[1],
         BlakeHash.f[0], BlakeHash.f[1]);
  printf(".buf = {");
  for (register uint8_t i = 0; i < BLAKE2B_BLOCKBYTES; i++)
    printf("%" PRIu8 ", ", BlakeHash.buf[i]);
  puts("\n");
  printf("}\n.buflen = %d\n.outlen = %d\n",
          BlakeHash.buflen, BlakeHash.outlen);
  printf(".last_node = %" PRIu8 "\n", BlakeHash.last_node);
  fflush(stdout);
}*/

static const blake2b_state miou = {
    .h = {
         UINT64_C(7640891576939301128), UINT64_C(13503953896175478587),
         UINT64_C(4354685564936845355), UINT64_C(11912009170470909681),
         UINT64_C(5840696475078001361), UINT64_C(11170449401992604703),
         UINT64_C(2270897969802886507), UINT64_C(6620516959819538809)
    },
    .t = {UINT64_C(0), UINT64_C(0)},
    .f = {UINT64_C(0), UINT64_C(0)},
    .buf = {
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
    },
    .buflen = 0,
    .outlen = 64,
    .last_node = 0
};


int blake2b_init_param(blake2b_state *S, const blake2b_param *P) {
    const unsigned char *p = (const unsigned char *)P;
    unsigned int i;

    if (NULL == P || NULL == S) {
        return -1;
    }

    blake2b_init0(S);

    /* IV XOR Parameter Block */
    for (i = 0; i < 8; ++i) {
        S->h[i] ^= load64(&p[i * sizeof(S->h[i])]);
    }
    S->outlen = P->digest_length;
    return 0;
}

void compare_buffs(uint64_t *h, size_t outlen)
{
    // printf("CMP : %d", memcmp(h, miou.h, 8*(sizeof(uint64_t))));
    printf("miou : %" PRIu64 " - h : %" PRIu64 " - outlen : %ld\n", miou.h[0], h[0], outlen);
    fflush(stdout);
}

/* Sequential blake2b initialization */
int blake2b_init(blake2b_state *S, size_t outlen) {
    memcpy(S, &miou, sizeof(*S));
    S->h[0] += outlen;
    return 0;
}


void print64(const char *name, const uint64_t *array, uint16_t size) {
  printf("%s = {", name);
  for (uint8_t i = 0; i < size; i++) printf("UINT64_C(%" PRIu64 "), ", array[i]);
  printf("};\n");
}
int blake2b_init_key(blake2b_state *S, size_t outlen, const void *key,
                     size_t keylen) {
    return 0;
}

static void blake2b_compress(blake2b_state *S, const uint8_t *block) {
    uint64_t m[16];
    uint64_t v[16];
    unsigned int i, r;

    for (i = 0; i < 16; ++i) {
        m[i] = load64(block + i * 8);
    }

    for (i = 0; i < 8; ++i) {
        v[i] = S->h[i];
    }

    v[8] = blake2b_IV[0];
    v[9] = blake2b_IV[1];
    v[10] = blake2b_IV[2];
    v[11] = blake2b_IV[3];
    v[12] = blake2b_IV[4] ^ S->t[0];
    v[13] = blake2b_IV[5]/* ^ S->t[1]*/;
    v[14] = blake2b_IV[6] ^ S->f[0];
    v[15] = blake2b_IV[7]/* ^ S->f[1]*/;

#define G(r, i, a, b, c, d)                                                    \
    do {                                                                       \
        a = a + b + m[blake2b_sigma[r][2 * i + 0]];                            \
        d = rotr64(d ^ a, 32);                                                 \
        c = c + d;                                                             \
        b = rotr64(b ^ c, 24);                                                 \
        a = a + b + m[blake2b_sigma[r][2 * i + 1]];                            \
        d = rotr64(d ^ a, 16);                                                 \
        c = c + d;                                                             \
        b = rotr64(b ^ c, 63);                                                 \
    } while ((void)0, 0)

#define ROUND(r)                                                               \
    do {                                                                       \
        G(r, 0, v[0], v[4], v[8], v[12]);                                      \
        G(r, 1, v[1], v[5], v[9], v[13]);                                      \
        G(r, 2, v[2], v[6], v[10], v[14]);                                     \
        G(r, 3, v[3], v[7], v[11], v[15]);                                     \
        G(r, 4, v[0], v[5], v[10], v[15]);                                     \
        G(r, 5, v[1], v[6], v[11], v[12]);                                     \
        G(r, 6, v[2], v[7], v[8], v[13]);                                      \
        G(r, 7, v[3], v[4], v[9], v[14]);                                      \
    } while ((void)0, 0)

    for (r = 0; r < 12; ++r) ROUND(r);

    for (i = 0; i < 8; ++i) S->h[i] = S->h[i] ^ v[i] ^ v[i + 8];

#undef G
#undef ROUND
}

int blake2b_update(blake2b_state *S, const void *in, size_t inlen) {
    const uint8_t *pin = (const uint8_t *)in;
    /* Complete current block */
    memcpy(&S->buf[4], pin, 124);
    blake2b_increment_counter(S, BLAKE2B_BLOCKBYTES);
    blake2b_compress(S, S->buf);
    S->buflen = 0;
    pin += 124;

    register int8_t i = 7;
    /* Avoid buffer copies when possible */
    while (i--) {
      blake2b_increment_counter(S, BLAKE2B_BLOCKBYTES);
      blake2b_compress(S, pin);
      pin += BLAKE2B_BLOCKBYTES;
    }
    memcpy(&S->buf[S->buflen], pin, 4);
    S->buflen += 4;
    return 0;
}

void my_blake2b_update(blake2b_state *S, const void *in, size_t inlen) {

    memcpy(&S->buf[S->buflen], in, inlen);
    S->buflen += (unsigned int)inlen;
}

int blake2b_final(blake2b_state *S, void *out, size_t outlen) {
    uint8_t buffer[BLAKE2B_OUTBYTES] = {0};
    unsigned int i;

    blake2b_increment_counter(S, S->buflen);
    blake2b_set_lastblock(S);
    memset(&S->buf[S->buflen], 0, BLAKE2B_BLOCKBYTES - S->buflen); /* Padding */
    blake2b_compress(S, S->buf);

    for (i = 0; i < 8; ++i) { /* Output full hash to temp buffer */
        store64(buffer + sizeof(S->h[i]) * i, S->h[i]);
    }

    memcpy(out, buffer, S->outlen);

    burn(buffer, sizeof(buffer));
    burn(S->buf, sizeof(S->buf));
    burn(S->h, sizeof(S->h));
    return 0;
}

int blake2b(void *out, const void *in, const void *key, size_t keylen)
{
    blake2b_state S;

    blake2b_init(&S, 64);
    my_blake2b_update(&S, in, 64);
    blake2b_final(&S, out, 64);
    burn(&S, sizeof(S));
    return 0;
}

void blake2b_too(void *pout, const void *in)
{
    uint8_t *out = (uint8_t *)pout;
    uint8_t out_buffer[64];
    uint8_t in_buffer[64];

    blake2b_state blake_state;
    blake2b_init(&blake_state, 64);
    blake_state.buflen = blake_state.buf[1] = 4;
    my_blake2b_update(&blake_state, in, 72);
    blake2b_final(&blake_state, out_buffer, 64);
    memcpy(out, out_buffer, 32);
    out += 32;

    register uint8_t i = 29;
    while (i--) {
        memcpy(in_buffer, out_buffer, 64);
        blake2b(out_buffer, in_buffer, NULL, 0);
        memcpy(out, out_buffer, 32);
        out += 32;
    }

    memcpy(in_buffer, out_buffer, 64);
    blake2b(out_buffer, in_buffer, NULL, 0);
    memcpy(out, out_buffer, 64);

    burn(&blake_state, sizeof(blake_state));
}

/* Argon2 Team - Begin Code */
int blake2b_long(void *pout, const void *in)
{
    uint8_t *out = (uint8_t *)pout;
    blake2b_state blake_state;
    uint8_t outlen_bytes[sizeof(uint32_t)] = {0};

    store32(outlen_bytes, 32);

    blake2b_init(&blake_state, 32);
    my_blake2b_update(&blake_state, outlen_bytes, sizeof(outlen_bytes));
    blake2b_update(&blake_state, in, 1024);
    blake2b_final(&blake_state, out, 32);
    burn(&blake_state, sizeof(blake_state));
    return 0;

}
/* Argon2 Team - End Code */
