#include <iostream>
#include <array>
#include <cstdio>
#include <tuple>
#include <vector>
#include <string>

using namespace std;

typedef unsigned int word;
typedef unsigned char byte;

const word ONE = 1;
const word BYTE_MASK = 0xff;

word rotate(word num, int count) {
    return (num << count) | (num >> (32 - count));
}

void quarterround(word& y0, word& y1, word& y2, word& y3) {
    word z1 = y1 ^ rotate(y0 + y3, 7);
    word z2 = y2 ^ rotate(z1 + y0, 9);
    word z3 = y3 ^ rotate(z2 + z1, 13);
    word z0 = y0 ^ rotate(z3 + z2, 18);

    y0 = z0, y1 = z1, y2 = z2, y3 = z3;
}

void rowround(array <word, 16>& m) {
    quarterround(m[0], m[1], m[2], m[3]);
    quarterround(m[5], m[6], m[7], m[4]);
    quarterround(m[10], m[11], m[8], m[9]);
    quarterround(m[15], m[12], m[13], m[14]);
}

void columnround(array <word, 16>& m) {
    quarterround(m[0], m[4], m[8], m[12]);
    quarterround(m[5], m[9], m[13], m[1]);
    quarterround(m[10], m[14], m[2], m[6]);
    quarterround(m[15], m[3], m[7], m[11]);
}

void doubleround(array <word, 16>& m) {
    columnround(m);
    rowround(m);
}

word littleendian(byte b0, byte b1, byte b2, byte b3) {
    return b0 + (ONE << 8) * b1 + (ONE << 16) * b2 + (ONE << 24) * b3;
}

// 0 & 0 == 0
// 0 & 1 == 0
// 1 & 0 == 0
// 1 & 1 == 1
tuple <byte, byte, byte, byte> littleendian_inv(word n) {
    byte b0 = (n >> 0)  & BYTE_MASK;
    byte b1 = (n >> 8)  & BYTE_MASK;
    byte b2 = (n >> 16) & BYTE_MASK;
    byte b3 = (n >> 24) & BYTE_MASK;

    return make_tuple(b0, b1, b2, b3);
}

// The Salsa20 hash function
void hash_func(array <byte, 64>& input) {
    array <word, 16> seq;
    for (int i = 0; i < 16; i++) {
        int idx = i * 4;
        seq[i] = littleendian(input[idx], input[idx + 1], input[idx + 2], input[idx + 3]);
    }

    auto backup{seq};

    for (int i = 0; i < 10; i++) {
        doubleround(seq);
    }

    for (int i = 0; i < 16; i++) {
        // Overflow == mod 2^32, no UB
        seq[i] += backup[i];
    }

    for (int i = 0; i < 16; i++) {
        int idx = i * 4;
        tie(input[idx], input[idx + 1], input[idx + 2], input[idx + 3]) = littleendian_inv(seq[i]);
    }
}

// The Salsa20 expansion function for 32 byte key
array <byte, 64> expansion_32(const array <byte, 32>& key, const array <byte, 16>& n) {
    const static array <byte, 4> sigma0{101, 120, 112,  97};
    const static array <byte, 4> sigma1{110, 100,  32,  51};
    const static array <byte, 4> sigma2{ 50,  45,  98, 121};
    const static array <byte, 4> sigma3{116, 101,  32, 107};

    array <byte, 64> result;
    auto ptr = result.begin();

    copy(sigma0.begin(), sigma0.end(), ptr);
    advance(ptr, sigma0.size());

    copy(key.begin(), key.begin() + 16, ptr);
    advance(ptr, 16);

    copy(sigma1.begin(), sigma1.end(), ptr);
    advance(ptr, sigma1.size());

    copy(n.begin(), n.end(), ptr);
    advance(ptr, n.size());

    copy(sigma2.begin(), sigma2.end(), ptr);
    advance(ptr, sigma2.size());

    copy(key.begin() + 16, key.end(), ptr);
    advance(ptr, 16);

    copy(sigma3.begin(), sigma3.end(), ptr);

    hash_func(result);
    return result;
}

// The Salsa20 encryption function
vector <byte> encrypt(const vector <byte>& input, const array <byte, 32>& key, const array <byte, 8>& nonce) {
    array <byte, 16> to_expand;
    for (size_t i = 0; i < nonce.size(); i++) {
        to_expand[i] = nonce[i];
    }

    vector <byte> result(input.size());
    array <byte, 64> expanded;
    long long val = 0;
    for (size_t i = 0; i < input.size(); i++) {
        if (i % 64 == 0) {
            long long val_copy = val++;
            for (size_t j = 8; j < to_expand.size(); j++) {
                to_expand[j] = val_copy & BYTE_MASK;
                val_copy >>= 8;
            }

            expanded = expansion_32(key, to_expand);
        }

        result[i] = input[i] ^ expanded[i % 64];
    }

    return result;
}

int main() {
    string hello_world = "asdasdasdsa";

    cout << "Original message" << endl;
    cout << hello_world << endl;

    vector <byte> source(hello_world.size());
    for (size_t i = 0; i < source.size(); i++) {
        source[i] = hello_world[i];
    }

    array <byte, 32> key;
    for (int i = 0; i < 32; i++) {
        key[i] = 100 + i;
    }

    array <byte, 8> nonce;
    for (int i = 0; i < 8; i++) {
        nonce[i] = 200 + i;
    }

    vector <byte> encrypted(encrypt(source, key, nonce));
    vector <byte> decrypted(encrypt(encrypted, key, nonce));

    cout << "Source" << endl;
    for (byte x : source) {
        cout << (int)x << ' ';
    }
    cout << endl;

    cout << "Encrypted" << endl;
    for (byte x : encrypted) {
        cout << (int)x << ' ';
    }
    cout << endl;

    cout << "Decrypted" << endl;
    for (byte x : decrypted) {
        cout << (int)x << ' ';
    }
    cout << endl;

    return 0;
}