#pragma once
#include <array>
#include <algorithm>
#include <cstdint>

#define ROTL(a,b) (((a) << (b)) | ((a) >> (32 - (b))))
#define QR(a, b, c, d) (			\
	a += b,  d ^= a,  d = ROTL(d,16),	\
	c += d,  b ^= c,  b = ROTL(b,12),	\
	a += b,  d ^= a,  d = ROTL(d, 8),	\
	c += d,  b ^= c,  b = ROTL(b, 7))

//
// Based on code from RFC 7539
//
template<int ROUNDS>
void chacha_do_rounds(std::array<uint32_t, 16>& out) {
	static_assert(ROUNDS%2 == 0, "ChaPRNG: Rounds count MUST be even.");
	for (int i = 0; i < ROUNDS; i += 2) {
		// Odd round
		QR(out[0], out[4], out[ 8], out[12]); // column 0
		QR(out[1], out[5], out[ 9], out[13]); // column 1
		QR(out[2], out[6], out[10], out[14]); // column 2
		QR(out[3], out[7], out[11], out[15]); // column 3
		// Even round
		QR(out[0], out[5], out[10], out[15]); // diagonal 1 (main diagonal)
		QR(out[1], out[6], out[11], out[12]); // diagonal 2
		QR(out[2], out[7], out[ 8], out[13]); // diagonal 3
		QR(out[3], out[4], out[ 9], out[14]); // diagonal 4
	}
}

template<int ROUNDS, bool REMIX = true, bool DIRECT = false>
std::array<uint32_t, 16> chacha_block(std::array<uint32_t, 16>& in)
{
	std::array<uint32_t, 16> out;
	int i;

	if constexpr(DIRECT) {
		chacha_do_rounds<ROUNDS>(in);
		return out;
	}else {
		for (i = 0; i < 16; ++i)
			out[i] = in[i];
		
		chacha_do_rounds<ROUNDS>(out);
		
		if constexpr(REMIX) {
			for (i = 0; i < 16; ++i)
				out[i] += in[i];
		}
		
		return out;
	}
}


template<int ROUNDS, bool REMIX = true, bool DIRECT = false>
class ChaPRNG {
	std::array<uint32_t, 16> source;
	std::array<uint32_t, 16> current;
	uint32_t it = 0;
public:
	ChaPRNG(std::array<uint32_t, 12> seed = {0}) {
		source[0] = 'expa';
		source[1] = 'nd 3';
		source[2] = '2-by';
		source[3] = 'te k';
		std::copy(seed.begin(), seed.end(), source.begin()+4);
	}
	
	uint32_t operator()() {
		auto spot = it & 0b1111;
		if(spot == 0) {
			source[12]++;
			if constexpr (DIRECT) {
				chacha_block<ROUNDS, REMIX, DIRECT>(source);
			} else {
				current = chacha_block<ROUNDS, REMIX, DIRECT>(source);
			}
		}
		++it;
		if constexpr (DIRECT)
			return source[spot];
		else
			return current[spot];
	}
};