package net.y3n20u.aeszip;

import static net.y3n20u.aeszip.CommonValues.BLOCK_SIZE;

import java.io.FilterOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.NoSuchAlgorithmException;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

public class AesCtrBlockCipherOutputStream extends FilterOutputStream {

	// TODO: I don't know the right setting for 'padding'
	/** name of algorithm used for encryption */
	public static final String CIPHER_MODE_AES_CTR = "AES/CTR/NoPadding";

	/** name of algorithm used for key construction */
	public static final String KEY_ALGORITHM = "AES";

	public static final byte[] INITIAL_IV = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };

	private static final int NONCE_SIZE = 8;

	private static final byte BYTE_HEX_FF = (byte) 0xff;

	private static final String MESSAGE_NONCE_OVERFLOW = "nonce overflow (it must not occur).";

	private Cipher cipher;
	private Key encryptKey;
	private final byte[] iv = new byte[16];
	private final byte[] restBytes = new byte[16];
	private int restBytesLength;
	private final int mode;

	public static AesCtrBlockCipherOutputStream generateEncryptInstance(OutputStream out) {
		if (out == null) {
			// TODO
			throw new IllegalArgumentException();
		}
		return new AesCtrBlockCipherOutputStream(out, Cipher.ENCRYPT_MODE);
	}
	
	public static AesCtrBlockCipherOutputStream generateDecryptInstance(OutputStream out) {
		if (out == null) {
			// TODO
			throw new IllegalArgumentException();
		}
		return new AesCtrBlockCipherOutputStream(out, Cipher.DECRYPT_MODE);
	}
	
	private AesCtrBlockCipherOutputStream(OutputStream out, int mode) {
		super(out);
		if (Cipher.DECRYPT_MODE != mode && Cipher.ENCRYPT_MODE != mode) {
			// TODO
			throw new IllegalArgumentException();
		}
		this.mode = mode;
		try {
			cipher = Cipher.getInstance(CIPHER_MODE_AES_CTR);
		} catch (NoSuchAlgorithmException nsae) {
			throw new AesZipRuntimeException(nsae);
		} catch (NoSuchPaddingException nspe) {
			throw new AesZipRuntimeException(nspe);
		}
	}

	public void init(byte[] keyBytes) {
		System.arraycopy(INITIAL_IV, 0, iv, 0, BLOCK_SIZE);
		encryptKey = new SecretKeySpec(keyBytes, KEY_ALGORITHM);
		restBytesLength = 0;
	}

	private void initIvAndCipher() {
		IvParameterSpec ivParameter = new IvParameterSpec(iv);
		try {
			cipher.init(mode, encryptKey, ivParameter);
		} catch (InvalidKeyException ike) {
			throw new AesZipRuntimeException(ike);
		} catch (InvalidAlgorithmParameterException iape) {
			throw new AesZipRuntimeException(iape);
		}
	}

	@Override
	public void write(byte[] b, int off, int len) throws IOException {
		if (restBytesLength > 0) {
			int length = BLOCK_SIZE - restBytesLength;
			if (len < length) {
				System.arraycopy(b, off, restBytes, restBytesLength, len);
				restBytesLength += len;
				return;
			}
			byte[] buffer = new byte[BLOCK_SIZE];
			System.arraycopy(restBytes, 0, buffer, 0, restBytesLength);
			System.arraycopy(b, off, buffer, restBytesLength, length);
			encryptAndWriteBlock(buffer, 0, BLOCK_SIZE);
			off += length;
			len -= length;
			restBytesLength = 0;
		}
		int i = off;
		for (; i + BLOCK_SIZE <= off + len; i += BLOCK_SIZE) {
			encryptAndWriteBlock(b, i, BLOCK_SIZE);
		}
		int m = off + len - i;
		if (m > 0) {
			restBytesLength = m;
			System.arraycopy(b, i, restBytes, 0, m);
		}
	}

	private void encryptAndWriteBlock(byte[] b, int off, int len) throws IOException {
		try {
			this.incrementIv();
			this.initIvAndCipher();
			byte[] buffer = cipher.doFinal(b, off, len);
			this.out.write(buffer, 0, buffer.length);
		} catch (IllegalBlockSizeException ibse) {
			throw new AesZipRuntimeException(ibse);
		} catch (BadPaddingException bpe) {
			throw new AesZipRuntimeException(bpe);
		}
	}

	private void incrementIv() {
		for (int i = 0; i < NONCE_SIZE; i++) {
			if (iv[i] != BYTE_HEX_FF) {
				iv[i]++;
				return;
			}
			iv[i] = 0;
		}
		throw new AesZipRuntimeException(MESSAGE_NONCE_OVERFLOW);
	}

	@Override
	public void flush() throws IOException {
		if (restBytesLength > 0) {
			encryptAndWriteBlock(restBytes, 0, restBytesLength);
			restBytesLength = 0;
		}
		this.out.flush();
	}

}
