#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <immintrin.h>
#include <math.h>

struct md5x8_context {
    __m256i k[64];
    __m256i iv[4];

    __m256i a, b, c, d;
};

void md5x8_pre_init(struct md5x8_context* ctx) {
    for (uint32_t i = 0; i < 64; ++i) {
        uint32_t x = floor(fabs(sin(i + 1)) * (double)0x100000000ull);
        ctx->k[i] = _mm256_set1_epi32(x);
    }
    ctx->iv[0] = _mm256_set1_epi32(0x67452301);
    ctx->iv[1] = _mm256_set1_epi32(0xefcdab89);
    ctx->iv[2] = _mm256_set1_epi32(0x98badcfe);
    ctx->iv[3] = _mm256_set1_epi32(0x10325476);
}

void md5x8_init(struct md5x8_context* ctx) {
    ctx->a = ctx->iv[0];
    ctx->b = ctx->iv[1];
    ctx->c = ctx->iv[2];
    ctx->d = ctx->iv[3];
}

inline __m256i md5x8_fx(int s, __m256i a, __m256i b, __m256i k, __m256i x, __m256i f) {
    f = _mm256_add_epi32(_mm256_add_epi32(f, a), _mm256_add_epi32(k, x));
    return _mm256_add_epi32(b, _mm256_or_si256(_mm256_slli_epi32(f, s), _mm256_srli_epi32(f, 32 - s)));
}

inline __m256i md5x8_f1(int s, __m256i a, __m256i b, __m256i c, __m256i d, __m256i k, __m256i x) {
    return md5x8_fx(s, a, b, k, x, _mm256_or_si256(_mm256_and_si256(b, c), _mm256_andnot_si256(b, d)));
}

inline __m256i md5x8_f2(int s, __m256i a, __m256i b, __m256i c, __m256i d, __m256i k, __m256i x) {
    return md5x8_fx(s, a, b, k, x, _mm256_or_si256(_mm256_and_si256(d, b), _mm256_andnot_si256(d, c)));
}

inline __m256i md5x8_f3(int s, __m256i a, __m256i b, __m256i c, __m256i d, __m256i k, __m256i x) {
    return md5x8_fx(s, a, b, k, x, _mm256_xor_si256(_mm256_xor_si256(b, c), d));
}

inline __m256i md5x8_f4(int s, __m256i a, __m256i b, __m256i c, __m256i d, __m256i k, __m256i x) {
    return md5x8_fx(s, a, b, k, x, _mm256_xor_si256(c, _mm256_or_si256(b, _mm256_xor_si256(d, _mm256_set1_epi32(0xFFFFFFFF)))));
}

void md5x8_raw_update(struct md5x8_context* ctx, const uint8_t blocks[8][64])
{
    __m256i x[16] = {0};
    for (size_t i = 0; i < 2; ++i) {
        __m256i x0 = _mm256_loadu_si256((const __m256i*)blocks[0] + i);
        __m256i x1 = _mm256_loadu_si256((const __m256i*)blocks[1] + i);
        __m256i x2 = _mm256_loadu_si256((const __m256i*)blocks[2] + i);
        __m256i x3 = _mm256_loadu_si256((const __m256i*)blocks[3] + i);
        __m256i x4 = _mm256_loadu_si256((const __m256i*)blocks[4] + i);
        __m256i x5 = _mm256_loadu_si256((const __m256i*)blocks[5] + i);
        __m256i x6 = _mm256_loadu_si256((const __m256i*)blocks[6] + i);
        __m256i x7 = _mm256_loadu_si256((const __m256i*)blocks[7] + i);

        __m256i t0 = _mm256_unpacklo_epi32(x0, x1);
        __m256i t1 = _mm256_unpackhi_epi32(x0, x1);
        __m256i t2 = _mm256_unpacklo_epi32(x2, x3);
        __m256i t3 = _mm256_unpackhi_epi32(x2, x3);
        __m256i t4 = _mm256_unpacklo_epi32(x4, x5);
        __m256i t5 = _mm256_unpackhi_epi32(x4, x5);
        __m256i t6 = _mm256_unpacklo_epi32(x6, x7);
        __m256i t7 = _mm256_unpackhi_epi32(x6, x7);

        __m256i s0 = _mm256_unpacklo_epi64(t0, t2);
        __m256i s1 = _mm256_unpackhi_epi64(t0, t2);
        __m256i s2 = _mm256_unpacklo_epi64(t1, t3);
        __m256i s3 = _mm256_unpackhi_epi64(t1, t3);
        __m256i s4 = _mm256_unpacklo_epi64(t4, t6);
        __m256i s5 = _mm256_unpackhi_epi64(t4, t6);
        __m256i s6 = _mm256_unpacklo_epi64(t5, t7);
        __m256i s7 = _mm256_unpackhi_epi64(t5, t7);

        x[i * 8 + 0] = _mm256_permute2x128_si256(s0, s4, 0x20);
        x[i * 8 + 1] = _mm256_permute2x128_si256(s1, s5, 0x20);
        x[i * 8 + 2] = _mm256_permute2x128_si256(s2, s6, 0x20);
        x[i * 8 + 3] = _mm256_permute2x128_si256(s3, s7, 0x20);
        x[i * 8 + 4] = _mm256_permute2x128_si256(s0, s4, 0x31);
        x[i * 8 + 5] = _mm256_permute2x128_si256(s1, s5, 0x31);
        x[i * 8 + 6] = _mm256_permute2x128_si256(s2, s6, 0x31);
        x[i * 8 + 7] = _mm256_permute2x128_si256(s3, s7, 0x31);
    }

    __m256i a = ctx->a;
    __m256i b = ctx->b;
    __m256i c = ctx->c;
    __m256i d = ctx->d;

    a = md5x8_f1(7, a, b, c, d, ctx->k[0], x[0]);
    d = md5x8_f1(12, d, a, b, c, ctx->k[1], x[1]);
    c = md5x8_f1(17, c, d, a, b, ctx->k[2], x[2]);
    b = md5x8_f1(22, b, c, d, a, ctx->k[3], x[3]);
    a = md5x8_f1(7, a, b, c, d, ctx->k[4], x[4]);
    d = md5x8_f1(12, d, a, b, c, ctx->k[5], x[5]);
    c = md5x8_f1(17, c, d, a, b, ctx->k[6], x[6]);
    b = md5x8_f1(22, b, c, d, a, ctx->k[7], x[7]);
    a = md5x8_f1(7, a, b, c, d, ctx->k[8], x[8]);
    d = md5x8_f1(12, d, a, b, c, ctx->k[9], x[9]);
    c = md5x8_f1(17, c, d, a, b, ctx->k[10], x[10]);
    b = md5x8_f1(22, b, c, d, a, ctx->k[11], x[11]);
    a = md5x8_f1(7, a, b, c, d, ctx->k[12], x[12]);
    d = md5x8_f1(12, d, a, b, c, ctx->k[13], x[13]);
    c = md5x8_f1(17, c, d, a, b, ctx->k[14], x[14]);
    b = md5x8_f1(22, b, c, d, a, ctx->k[15], x[15]);

    a = md5x8_f2(5, a, b, c, d, ctx->k[16], x[1]);
    d = md5x8_f2(9, d, a, b, c, ctx->k[17], x[6]);
    c = md5x8_f2(14, c, d, a, b, ctx->k[18], x[11]);
    b = md5x8_f2(20, b, c, d, a, ctx->k[19], x[0]);
    a = md5x8_f2(5, a, b, c, d, ctx->k[20], x[5]);
    d = md5x8_f2(9, d, a, b, c, ctx->k[21], x[10]);
    c = md5x8_f2(14, c, d, a, b, ctx->k[22], x[15]);
    b = md5x8_f2(20, b, c, d, a, ctx->k[23], x[4]);
    a = md5x8_f2(5, a, b, c, d, ctx->k[24], x[9]);
    d = md5x8_f2(9, d, a, b, c, ctx->k[25], x[14]);
    c = md5x8_f2(14, c, d, a, b, ctx->k[26], x[3]);
    b = md5x8_f2(20, b, c, d, a, ctx->k[27], x[8]);
    a = md5x8_f2(5, a, b, c, d, ctx->k[28], x[13]);
    d = md5x8_f2(9, d, a, b, c, ctx->k[29], x[2]);
    c = md5x8_f2(14, c, d, a, b, ctx->k[30], x[7]);
    b = md5x8_f2(20, b, c, d, a, ctx->k[31], x[12]);

    a = md5x8_f3(4, a, b, c, d, ctx->k[32], x[5]);
    d = md5x8_f3(11, d, a, b, c, ctx->k[33], x[8]);
    c = md5x8_f3(16, c, d, a, b, ctx->k[34], x[11]);
    b = md5x8_f3(23, b, c, d, a, ctx->k[35], x[14]);
    a = md5x8_f3(4, a, b, c, d, ctx->k[36], x[1]);
    d = md5x8_f3(11, d, a, b, c, ctx->k[37], x[4]);
    c = md5x8_f3(16, c, d, a, b, ctx->k[38], x[7]);
    b = md5x8_f3(23, b, c, d, a, ctx->k[39], x[10]);
    a = md5x8_f3(4, a, b, c, d, ctx->k[40], x[13]);
    d = md5x8_f3(11, d, a, b, c, ctx->k[41], x[0]);
    c = md5x8_f3(16, c, d, a, b, ctx->k[42], x[3]);
    b = md5x8_f3(23, b, c, d, a, ctx->k[43], x[6]);
    a = md5x8_f3(4, a, b, c, d, ctx->k[44], x[9]);
    d = md5x8_f3(11, d, a, b, c, ctx->k[45], x[12]);
    c = md5x8_f3(16, c, d, a, b, ctx->k[46], x[15]);
    b = md5x8_f3(23, b, c, d, a, ctx->k[47], x[2]);

    a = md5x8_f4(6, a, b, c, d, ctx->k[48], x[0]);
    d = md5x8_f4(10, d, a, b, c, ctx->k[49], x[7]);
    c = md5x8_f4(15, c, d, a, b, ctx->k[50], x[14]);
    b = md5x8_f4(21, b, c, d, a, ctx->k[51], x[5]);
    a = md5x8_f4(6, a, b, c, d, ctx->k[52], x[12]);
    d = md5x8_f4(10, d, a, b, c, ctx->k[53], x[3]);
    c = md5x8_f4(15, c, d, a, b, ctx->k[54], x[10]);
    b = md5x8_f4(21, b, c, d, a, ctx->k[55], x[1]);
    a = md5x8_f4(6, a, b, c, d, ctx->k[56], x[8]);
    d = md5x8_f4(10, d, a, b, c, ctx->k[57], x[15]);
    c = md5x8_f4(15, c, d, a, b, ctx->k[58], x[6]);
    b = md5x8_f4(21, b, c, d, a, ctx->k[59], x[13]);
    a = md5x8_f4(6, a, b, c, d, ctx->k[60], x[4]);
    d = md5x8_f4(10, d, a, b, c, ctx->k[61], x[11]);
    c = md5x8_f4(15, c, d, a, b, ctx->k[62], x[2]);
    b = md5x8_f4(21, b, c, d, a, ctx->k[63], x[9]);

    ctx->a = _mm256_add_epi32(ctx->a, a);
    ctx->b = _mm256_add_epi32(ctx->b, b);
    ctx->c = _mm256_add_epi32(ctx->c, c);
    ctx->d = _mm256_add_epi32(ctx->d, d);
}

void md5x8_final(struct md5x8_context* ctx, uint8_t out[8][32]) {
    __m256i x0 = ctx->a;
    __m256i x1 = ctx->b;
    __m256i x2 = ctx->c;
    __m256i x3 = ctx->d;
    __m256i x4 = _mm256_set1_epi32(0);
    __m256i x5 = _mm256_set1_epi32(0);
    __m256i x6 = _mm256_set1_epi32(0);
    __m256i x7 = _mm256_set1_epi32(0);

    __m256i t0 = _mm256_unpacklo_epi32(x0, x1);
    __m256i t1 = _mm256_unpackhi_epi32(x0, x1);
    __m256i t2 = _mm256_unpacklo_epi32(x2, x3);
    __m256i t3 = _mm256_unpackhi_epi32(x2, x3);
    __m256i t4 = _mm256_unpacklo_epi32(x4, x5);
    __m256i t5 = _mm256_unpackhi_epi32(x4, x5);
    __m256i t6 = _mm256_unpacklo_epi32(x6, x7);
    __m256i t7 = _mm256_unpackhi_epi32(x6, x7);

    __m256i s0 = _mm256_unpacklo_epi64(t0, t2);
    __m256i s1 = _mm256_unpackhi_epi64(t0, t2);
    __m256i s2 = _mm256_unpacklo_epi64(t1, t3);
    __m256i s3 = _mm256_unpackhi_epi64(t1, t3);
    __m256i s4 = _mm256_unpacklo_epi64(t4, t6);
    __m256i s5 = _mm256_unpackhi_epi64(t4, t6);
    __m256i s6 = _mm256_unpacklo_epi64(t5, t7);
    __m256i s7 = _mm256_unpackhi_epi64(t5, t7);

    _mm256_storeu_si256((__m256i*)(out + 0), _mm256_permute2x128_si256(s0, s4, 0x20));
    _mm256_storeu_si256((__m256i*)(out + 1), _mm256_permute2x128_si256(s1, s5, 0x20));
    _mm256_storeu_si256((__m256i*)(out + 2), _mm256_permute2x128_si256(s2, s6, 0x20));
    _mm256_storeu_si256((__m256i*)(out + 3), _mm256_permute2x128_si256(s3, s7, 0x20));
    _mm256_storeu_si256((__m256i*)(out + 4), _mm256_permute2x128_si256(s0, s4, 0x31));
    _mm256_storeu_si256((__m256i*)(out + 5), _mm256_permute2x128_si256(s1, s5, 0x31));
    _mm256_storeu_si256((__m256i*)(out + 6), _mm256_permute2x128_si256(s2, s6, 0x31));
    _mm256_storeu_si256((__m256i*)(out + 7), _mm256_permute2x128_si256(s3, s7, 0x31));
}


void unsafe_pad_block(uint8_t block[64], size_t length) {
    block[length] = 0x80;
    *(uint64_t*)(block + 56) = length * 8;
}

int main() {
    struct md5x8_context ctx;
    md5x8_pre_init(&ctx);

    uint8_t data[8][64] = {
        "We are the others",
        "We are the cast-outs",
        "We're the outsiders",
        "But you can't hide us",
    };

    unsafe_pad_block(data[0], 17);
    unsafe_pad_block(data[1], 20);
    unsafe_pad_block(data[2], 19);
    unsafe_pad_block(data[3], 21);

    uint8_t digests[8][32] = {0};

    md5x8_init(&ctx);
    md5x8_raw_update(&ctx, data);
    md5x8_final(&ctx, digests);

    for (size_t i = 0; i < 8; ++i) {
        for (size_t j = 0; j < 16; ++j)
            printf("%02X ", digests[i][j]);
        printf("\n");
    }

    return 0;
}