package aalg.crypto;

import java.util.Arrays;

import aalg.algebra.Utils;
import aalg.crypto.AesBase;
import aalg.crypto.IBlockCipher;


/**
 * Decryptor of AES block cipher.
 * @author Ing. Petr Ležák
 */
public class AesDecryptor extends AesBase implements IBlockCipherTransformation {
	
	private static final int[][] INV_MIX_COLUMN_MATRIX = {
		{14, 11, 13, 9},
		{9, 14, 11, 13},
		{13, 9, 14, 11},
		{11, 13, 9, 14}
	};

	private static final int[] INV_MIX_COLUMN_TABLE = createInvMixColumnTable();
	private static final int[] ROUND_TABLE = createRoundTable();

	private static int[] createInvMixColumnTable() {
		final int[] table = new int[BYTE_LAST * ROW_COUNT];

		for (int row = 0; row < ROW_COUNT; row++) {
			for (int value = BYTE_FIRST; value < BYTE_LAST; value++) {
				int result = 0;
				
				for (int targetRow = 0; targetRow < ROW_COUNT; targetRow++) {
					final int resultByte = multiplyBytes(value, INV_MIX_COLUMN_MATRIX[targetRow][row]);
					
					result |= resultByte << (Byte.SIZE * targetRow);
				}
				
				final int tableIndex = calculateRoundTableIndex(row, value);
				table[tableIndex] = result;
			}
		}
		
		return table;
	}

	private static int[] createRoundTable() {
		final int[] table = new int[BYTE_LAST * ROW_COUNT];
		
		for (int row = 0; row < ROW_COUNT; row++) {
			for (int value = BYTE_FIRST; value < BYTE_LAST; value++) {
				final int sub = subByteBackward(value);
				final int origTableIndex = calculateRoundTableIndex(row, value);
				final int subTableIndex = calculateRoundTableIndex(row, sub);
				
				table[origTableIndex] = INV_MIX_COLUMN_TABLE[subTableIndex];
			}
		}

		return table;
	}

	private final AesCipher cipher;
	private final int[] roundKeys = new int[MAX_ROUNDS * COLUMN_COUNT];
	
	
	public AesDecryptor(final AesCipher cipher) {
		this.cipher = cipher;
	}
	
	private void invSubBytes(final int[] block) {
		for (int i = 0; i < COLUMN_COUNT; i++) {
			final int origColumn = block[i];
			int transformedColumn = 0;
			
			for (int j = (ROW_COUNT - 1) * Byte.SIZE; j >= 0; j -= Byte.SIZE) {
				transformedColumn |= subByteBackward((origColumn >>> j) & 0xFF) << j;
			}
			
			block[i] = transformedColumn;
		}
	}

	private static int calculateRoundTableIndex (final int row, final int value) {
		return row * BYTE_LAST + value;
	}
	

	private void invShiftRows(final int[] block) {
		for (int i = 1; i < ROW_COUNT; i++) {
			final int shift = i * Byte.SIZE;
			final int mask = ~(0xFF << shift);
			int data = 0;
			
			for (int j = COLUMN_COUNT - 1; j >= 0; j--)
				data = (data << Byte.SIZE) | ((block[j] >>> shift) & 0xFF);
			
			data = Integer.rotateLeft(data, i * Byte.SIZE);
			
			for (int j = 0; j < COLUMN_COUNT; j++) {
				block[j] = (block[j] & mask) | ((data & 0xFF) << shift);
				data >>>= Byte.SIZE;
			}
		}
	}

	private void addRoundKey(final int round, final int[] block) {
		final int roundKeyAddress = ROW_COUNT * round;
		
		for (int i = 0; i < ROW_COUNT; i++) {
			block[i] ^= roundKeys[roundKeyAddress + i];
		}
	}
	
	public void oneRound(final int round, final int[] prevBlock, final int[] nextBlock) {
		Arrays.fill(nextBlock, 0);
		
		for (int columnIndex = 0; columnIndex < COLUMN_COUNT; columnIndex++) {
			int column = prevBlock[columnIndex];
			
			for (int rowIndex = 0; rowIndex < ROW_COUNT; rowIndex++) {
				final int value = column & 0xFF;
				final int targetColumnIndex = (columnIndex + rowIndex) & 0x03;
				final int tableIndex = calculateRoundTableIndex(rowIndex, value);
				
				nextBlock[targetColumnIndex] ^= ROUND_TABLE[tableIndex];
				column >>>= Byte.SIZE;
			}
		}
		
		addRoundKey(round, nextBlock);
	}

	@Override
	public IBlockCipher getCipher() {
		return cipher;
	}

	@Override
	public void transformBlock (final byte[] sourceBlock, final int sourceOffset, final byte[] targetBlock, final int targetOffset) {
		final int[] prevBlock = new int[ROW_COUNT];
		final int[] nextBlock = new int[ROW_COUNT];
		
		Utils.byteBlockToIntBlock(sourceBlock, prevBlock, sourceOffset, 0, COLUMN_COUNT);
		addRoundKey(roundCount, prevBlock);
		
		for (int round = roundCount - 1; round > 0; round--) {
			oneRound(round, prevBlock, nextBlock);
			System.arraycopy(nextBlock, 0, prevBlock, 0, ROW_COUNT);
		}
		
		invSubBytes(prevBlock);
		invShiftRows(prevBlock);
		addRoundKey(0, prevBlock);
		
		Utils.intBlockToByteBlock(prevBlock, targetBlock, 0, targetOffset, COLUMN_COUNT);
	}

	@Override
	public void setRoundKey(final int round, final int[] key, final int offset) {
		if (round == 0 || round == roundCount) {
			System.arraycopy(key, offset, roundKeys, round * COLUMN_COUNT, COLUMN_COUNT);
		}
		else {
			final int base = round * COLUMN_COUNT;
			
			for (int i = 0; i < COLUMN_COUNT; i++) {
				roundKeys[base + i] = invMixColumns (key[offset + i]);
			}
		}
	}
	
	public static int invMixColumns (final int column) {
		int result = 0;
		
		for (int row = 0; row < ROW_COUNT; row++) {
			final int shift = Byte.SIZE * row;
			final int origByte = (column >>> shift) & 0xFF;
			
			final int component = INV_MIX_COLUMN_TABLE[calculateRoundTableIndex(row, origByte)];
			
			final int component2 = invMixColumnsStandAlone(origByte << shift);
			
			if (component != component2)
				throw new RuntimeException("XXXXXXXXXX " + Long.toHexString(component & 0xFFFFFFFFL) + " " + Long.toHexString(component2 & 0xFFFFFFFFL));
			
			result ^= component;
		}
		
		return result;
	}
	
	public static int invMixColumnsStandAlone(final int column) {
		int result = 0;
		
		for (int resultRow = ROW_COUNT - 1; resultRow >= 0; resultRow--) {
			int resultByte = 0;
			
			for (int sourceRow = 0; sourceRow < ROW_COUNT; sourceRow++) {
				final int shift = Byte.SIZE * sourceRow;
				final int sourceByte = (column >>> shift) & 0xFF;
				
				resultByte ^= multiplyBytes(sourceByte, INV_MIX_COLUMN_MATRIX[resultRow][sourceRow]);
			}
			
			result = (result << Byte.SIZE) | resultByte;
		}
		
		return result;
	}

}
