// https://tools.ietf.org/html/rfc7539
import {MASK,int32s2bytes,bytes2int32s,zeroPad,createRandomNonce} from "./util.js";

function lrot(x,shift){
	return (x<<shift|x>>>(32-shift))&MASK;
}

/**
 * A Chacha20 cipher class.
 * @param {Array} key Array of bytes (integers: 0<=x<256). Short keys are padded to 32B, long keys are silently truncated.
 * @param {Array} nonce optional. If present, it must be an Array of bytes (integers: 0<=x<256). Short nonces are padded to 12B, long nonces are silently truncated.
 */
export function Chacha20(key,nonce){
	const NONCE_LEN=12;
	if(nonce===undefined){
		nonce=createRandomNonce(NONCE_LEN);
	}
	this._nonce=zeroPad(nonce,NONCE_LEN);
	nonce=bytes2int32s(this._nonce);
	key=bytes2int32s(zeroPad(key,32));
	
	this._state=[
		0x61707865,0x3320646e,0x79622d32,0x6b206574,
		key[0],key[1],key[2],key[3],
		key[4],key[5],key[6],key[7],
		1,nonce[0],nonce[1],nonce[2]
	];
	this._buffer=[];
}

Chacha20.prototype.setPos=function(pos){
	this._state[12]=pos;
	this._buffer=[];
}

Chacha20.prototype.getByte=function(){
	if(this._buffer.length==0){
		this._buffer=int32s2bytes(this._computeBlock());
		this._incrementPos();
	}
	return this._buffer.shift();
};

Chacha20.prototype.getNonce=function(){
	return this._nonce.concat();
};

Chacha20.prototype._quarterRound=function(arr,ia,ib,ic,id){
	let a=arr[ia], b=arr[ib], c=arr[ic], d=arr[id];
	a=(a+b)&MASK; d=lrot(d^a,16);
	c=(c+d)&MASK; b=lrot(b^c,12);
	a=(a+b)&MASK; d=lrot(d^a,8);
	c=(c+d)&MASK; b=lrot(b^c,7);
	arr[ia]=a; arr[ib]=b; arr[ic]=c; arr[id]=d;
};

Chacha20.prototype._computeBlock=function(){
	let state=this._state.slice();

	for(let i=0;i<10;i++){ // 10 double rounds
		// column round
		this._quarterRound(state,0,4,8,12);
		this._quarterRound(state,1,5,9,13);
		this._quarterRound(state,2,6,10,14);
		this._quarterRound(state,3,7,11,15);
		// diagonal round
		this._quarterRound(state,0,5,10,15);
		this._quarterRound(state,1,6,11,12);
		this._quarterRound(state,2,7,8,13);
		this._quarterRound(state,3,4,9,14);
	}

	return state.map((xi,i)=>(xi+this._state[i])&MASK);
};

Chacha20.prototype._incrementPos=function(){
	this._state[12]=(this._state[12]+1)&MASK;
};

export function encrypt(data,key,nonce){
	let cipher=new Chacha20(key,nonce);
	nonce=cipher.getNonce();
	return [nonce,data.map(b=>b^cipher.getByte())];
}

export function decrypt(ciphertext,key,nonce){
	let cipher=new Chacha20(key,nonce);
	return ciphertext.map(b=>b^cipher.getByte());
}