// 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());
}