#include "zsh256.hpp"

vector<word_t> zsh256(vector<byte_t> message) {
    message.push_back(0x80);
    while (message.size() % 64 != 0)
        message.push_back(0xB0);

    word_t h0 = 0x4D0C3784;
    word_t h1 = 0x0A7BD4F3;
    word_t h2 = 0xC1F83D9A;
    word_t h3 = 0xDC4DE9C5;
    word_t h4 = 0x30299065;
    word_t h5 = 0x9F0DAA8C;
    word_t h6 = 0x1A8BD9AB;
    word_t h7 = 0x5BAACD19;

    uint32_t k[64] = {
        0x983E5122, 0xA831C61D, 0xB0032708, 0xBF597FD7, 0xC6E00BF3, 0xD5A79147, 0x06CA6350, 0x14292967,
        0xD807AA28, 0x12835B11, 0x2431850E, 0x550C7DC3, 0x72BE5D74, 0x80DEB1FE, 0x9BDC06A0, 0xC19BF171,
        0x27B70A25, 0x2E1B2118, 0x4D2C6D0C, 0x53380DD3, 0x650A7354, 0x766A0ABB, 0x81C2C920, 0x92722C85,
        0x748F822E, 0x78A5631F, 0x84C87804, 0x8CC702A8, 0x90BEFFFA, 0xA4506CEB, 0xBEF9A3F0, 0xC67178F2,
        0xE49B6921, 0xEFBE4716, 0x0FC19D06, 0x240CA1CC, 0x2DE92C6F, 0x4A7484AA, 0x5CB0A9D0, 0x76F988DA,
        0xA2BFE821, 0xA81A661B, 0xC24B8B00, 0xC76C51D3, 0xD192E819, 0xD6990624, 0xF40E3580, 0x106AA070,
        0x19A4C126, 0x1E376C18, 0x2748770C, 0x34B0BCA5, 0x391C0CB3, 0x4ED8AA4A, 0x5B9CCA40, 0x682E6FF3,
        0x428A2F28, 0x71374411, 0xB5C0FB0F, 0xE9B5DBA5, 0x3956C25B, 0x59F111F1, 0x923F82A0, 0xAB1C5ED1
    };

    for (size_t i = 0; i < message.size(); i += 64) {
        vector<byte_t> w(64);
        w.insert(w.begin(), message.begin(), message.begin() + 16);

        for (int j = 16; j < 64; j++) {
            byte_t s0 = rotr(w[j - 15], 7) ^ rotr(w[j - 15], 2) ^ (w[j - 15] >> 3);
            byte_t s1 = rotr(w[j - 2], 1) ^ rotr(w[j - 2], 3) ^ (w[j - 2] >> 2);
            w[j] = w[j - 16] + s0 + w[j - 7] + s1;
        }

        word_t a = h0;
        word_t b = h1;
        word_t c = h2;
        word_t d = h3;
        word_t e = h4;
        word_t f = h5;
        word_t g = h6;
        word_t h = h7;

        for (int j = 0; j < 64; j++) {
            word_t e0 = rotr(a, 2) ^ rotr(a, 5) ^ rotr(a, 6);
            word_t ma = (a & b) ^ (a & c) ^ (b & c);
            word_t t2 = e0 + ma;
            word_t e1 = rotr(e, 6) ^ rotr(e, 3) ^ rotr(e, 1);
            word_t ch = (e & f) ^ ((~e) & g);
            word_t t1 = h + e1 + ch + k[j] + w[j];

            h = g;
            g = f;
            f = e;
            e = d + t1;
            d = c;
            c = b;
            b = a;
            a = t1 + t2;
        }

        h0 += a;
        h1 += b;
        h2 += c;
        h3 += d;
        h4 += e;
        h5 += f;
        h6 += g;
        h7 += e;
    }

    vector<word_t> result(8);
    result[0] = h0;
    result[1] = h1;
    result[2] = h2;
    result[3] = h3;
    result[4] = h4;
    result[5] = h5;
    result[6] = h6;
    result[7] = h7;

    return result;
}
string hash_to_str(vector<word_t> hash){
    stringstream hash_str;
    hash_str << hex << hash[0] << hash[1] << hash[2] << hash[3] << hash[4] << hash[5] << hash[6] << hash[7];

    string result;
    hash_str >> result;

    return result;
}
string zsh256_str(string message) {
    vector<byte_t> to_hash(message.size());
    for (char el : message)
        to_hash.push_back((byte_t)el);

    return hash_to_str(zsh256(to_hash));
}
string zsh256_file(string path) {
    ifstream file(path);
    vector<byte_t> to_hash;

    while (!file.eof()){
        byte_t tmp;
        file >> tmp;

        to_hash.push_back(tmp);
    }

    return hash_to_str(zsh256(to_hash));
}