#include <cstring>
#include <cassert>
#include "aes.h"

using namespace std;

namespace AES{

	//http://csrc.nist.gov/publications/fips/fips197/fips-197.pdf

	void subWord(uint8_t *word);

	void keyExpand(uint8_t *keysched,const uint8_t *key,int keylen,int numrounds);

	void addRoundKey(uint8_t *state,const uint8_t *roundkey);

	void subBytes(uint8_t *state){
		subWord(state);
		subWord(state+4);
		subWord(state+8);
		subWord(state+12);
	}

	void shiftRows(uint8_t *state);

	void mixColumns(uint8_t *state);


	void encryptBlock(uint8_t *state,const uint8_t *key,const uint8_t *data,int keylen,int numrounds){
		uint8_t keysched[16*(numrounds+1)];
		keyExpand(keysched,key,keylen,numrounds);
		memcpy(state,data,16);

		addRoundKey(state,keysched);
		for(int round=0;round<numrounds-1;round++){
			subBytes(state);
			shiftRows(state);
			mixColumns(state);
			addRoundKey(state,keysched+16*(round+1));
		}
		subBytes(state);
		shiftRows(state);
		addRoundKey(state,keysched+16*numrounds);
	}

	string encrypt(const string &key,const string &data,int numrounds){
		int sz=data.size();
		if(sz==0)return {};
		int blocks=sz/16;
		int padding=sz%16==0?16:16-sz%16;
		string res;
		assert((sz+padding)%16==0);
		res.reserve(sz+padding);
		uint8_t buf[16];
		for(int i=0;i<blocks;i++){
			encryptBlock(buf,(const uint8_t*)key.data(),(const uint8_t*)data.data()+16*i,key.size(),numrounds);
			res.insert(res.size(),(char*)buf,16);
		}
		uint8_t inbuf[16];
		if(padding<16)memcpy(inbuf,data.data()+16*blocks,16-padding);
		memset(inbuf+16-padding,padding,padding);
		encryptBlock(buf,(const uint8_t*)key.data(),inbuf,key.size(),numrounds);
		res.insert(res.size(),(char*)buf,16);
		return res;
	}

	string decrypt(const string &key,const string &data,int numrounds);

	string encrypt(const string &key,const string &data,Algorithm algo){
		assert(key.size()==4+2*algo);
		return encrypt(key,data,10+2*algo);
	}

	string decrypt(const string &key,const string &data,Algorithm algo)/*{
		assert(key.size()==4+2*algo);
		return decrypt(key,data,10+2*algo);
	}*/;

}