#pragma once
#include <stdint.h>
#include <iostream>
#include <stddef.h>
#include <gp/math.hpp>
#include <gp/array.hpp>
#include <gp/algorithm/rotate.hpp>
// BUG: A proper investigation is required to fix this file so that it gives the correct output.


template<typename word_t = uint32_t, size_t r = 20, size_t b = 128, word_t P = 0xb7e15163, word_t Q = 0x9e3779b9>
class RC6 {
	static constexpr size_t word_size = 8*sizeof(word_t);

	constexpr static word_t r_l(const word_t& w, size_t v) {
		return (w << v) | ( w >> (word_size-v));
	}

	constexpr static word_t r_r(const word_t& w, size_t v) {
		return (w >> v) | ( w << (word_size-v));
	}


	class RC6_KeySched {
		using sched_t = gp::array<word_t, 2*r+4>;
	public:
		static constexpr size_t c = (b+word_size-1)/word_size;
		static constexpr size_t v_3 = gp::max(c, 2*r+4);
		static constexpr size_t v = v_3*3;
	private:
		sched_t S{};
	public:
		constexpr RC6_KeySched(gp::array<word_t, c> L)
		{
			static_assert(r_l(r_r(13,13),13) == 13);

			auto it = S.begin();
			*(it++) = P;
			for(; it != S.end(); ++it)
			{
				*it = *(it-1) + Q;
			}
			word_t A = 0;
			word_t B = 0;
			word_t i = 0;
			word_t j = 0;
			
			size_t s = 0;
			while(true)
			{
				A = S[i] = r_l( S[i] + A + B, 3 );
				B += A;
				B = L[j] = r_l( L[j] + B, B%(word_size));
				++s;
				if(s >= v) break;
				i = s % S.size();
				j = s % L.size();
			}
		}
		

		const auto cbegin()
		{
			return S.cbegin();
		}

		const auto cend()
		{
			return S.cend();
		}

		const auto crbegin()
		{
			return S.crbegin();
		}

		const auto crend()
		{
			return S.crend();
		}
	};

	RC6_KeySched S;

public:

	typedef gp::array<word_t, RC6_KeySched::c> key_type;
	typedef gp::array<word_t, 4> block_type;

	constexpr RC6(const key_type& key)
	: S(key)
	{}

	constexpr block_type&& encrypt(block_type plaintext) {
		using namespace gp::math;
		auto& A = plaintext[0];
		auto& B = plaintext[1];
		auto& C = plaintext[2];
		auto& D = plaintext[3];

		auto it = S.cbegin();

		B += *(it++);
		D += *(it++);

		for(size_t i = 0; i < r; ++i)
		{
			auto u = r_l( D * ( 2 * D + 1 ), msb(word_size));
			auto t = r_l( B * ( 2 * B + 1 ), msb(word_size));
			A = r_l((A ^ t), u % word_size) + *(it++);
			C = r_l((C ^ u), t % word_size) + *(it++);
			gp::rotate(plaintext.begin(), plaintext.begin()+1, plaintext.end());
		}

		A += *(it++);
		C += *(it++);
		assert(it == S.cend());
		return std::move(plaintext);
	}

	constexpr block_type decrypt(block_type plaintext) {
		using namespace gp::math;
		auto& A = plaintext[0];
		auto& B = plaintext[1];
		auto& C = plaintext[2];
		auto& D = plaintext[3];
		auto it = S.crbegin();

		C -= *(it++);
		A -= *(it++);

		for(size_t i = 0; i < r; ++i)
		{
			gp::rotate(plaintext.begin(), plaintext.end()-1, plaintext.end());
			auto u = r_l( D * ( 2 * D + 1 ), msb(word_size));
			auto t = r_l( B * ( 2 * B + 1 ), msb(word_size));
			C = r_r( (C - *(it++)) , t % word_size) ^ u ;
			A = r_r( (A - *(it++)) , u % word_size) ^ t ;
		}

		D -= *(it++);
		B -= *(it++);
		assert(it == S.crend());
		return std::move(plaintext);
	}

};