package net.y3n20u.aeszip;

import static net.y3n20u.aeszip.CommonValues.BLOCK_SIZE;
import static net.y3n20u.aeszip.CommonValues.LENGTH_AUTHENTICATION_CODE;
import static net.y3n20u.aeszip.CommonValues.LENGTH_PASSWORD_VERIFICATION_VALUE;
import static net.y3n20u.aeszip.CommonValues.METHOD_DEFLATED;
import static net.y3n20u.aeszip.CommonValues.METHOD_STORED;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.io.RandomAccessFile;
import java.nio.charset.Charset;
import java.text.MessageFormat;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.zip.Inflater;
import java.util.zip.InflaterOutputStream;
import java.util.zip.ZipException;

import net.y3n20u.util.ByteHelper;

// FIXME: password check before generating the target file.
// FIXME: CRC or MAC check after generating files.
public class AesZipDecrypter {

	// this means 'read only.'
	private static final String FILE_ACCESS_MODE = "r";

	// skip 8 bytes of the 'End of central directory record.'
	// signature (4 bytes), number of this disk (2 bytes),
	// number of the disk with start of the central directory (2 bytes)
	private static final long OFFSET_NUMBER_OF_CENTRAL_DIR = 8;

	// private static final byte[] END_OF_CENTRAL_DIR_SIG = {0x05, 0x06, 0x4b,
	// 0x50};
	private static final byte[] END_OF_CENTRAL_DIR_SIG = { 0x50, 0x4b, 0x05, 0x06 };

	private static final byte[] CENTRAL_DIR_FILE_HEADER_SIG = { 0x50, 0x4b, 0x01, 0x02 };

	private static final byte[] AES_EXTRA_FIELD_HEADER_AND_SIZE = { 0x01, -0x67, 0x07, 0x00 };

	private static final byte[] AES_EXTRA_FIELD_VERSION_AE1 = { 0x01, 0x00 };
	private static final byte[] AES_EXTRA_FIELD_VERSION_AE2 = { 0x02, 0x00 };

	private static final byte[] AES_EXTRA_FIELD_VENDER_ID = { 0x41, 0x45 };

	private static final byte[] AES_ALLOWED_ENCRYPTION_STRENGTH = { 1, 2, 3 };

	private static final int LENGTH_EXTRA_FIELD = 11;

	private static final int ENCRYPTED_FLAG = 1;

	private static final int[] ALLOWED_ACTUAL_METHODS = { METHOD_STORED, METHOD_DEFLATED };

	private static final long LENGTH_LOCAL_FILE_HEADER_SKIP = 26;

	private final Map<String, AesZipEntry> nameToEntryMap;
	private final RandomAccessFile sourceFile;
	private final byte[] comment;

	// TODO: charset は entry に持たせる？
	private Charset fileNameCharset;
	private Charset commentCharset;
	private Charset passwordCharset;

	private final byte[] decryptionBlock = new byte[BLOCK_SIZE];

	// fields for progress listener.
	private ArchiveProgressListener progressListener;
	private ArchiveProgressStatus previousStatus;

	public AesZipDecrypter(File sourceFile) throws ZipException, IOException {
		if (sourceFile == null || sourceFile.isDirectory()) {
			// TODO
			throw new IllegalArgumentException();
		}
		this.sourceFile = new RandomAccessFile(sourceFile, FILE_ACCESS_MODE);
		fileNameCharset = Charset.defaultCharset();
		commentCharset = Charset.defaultCharset();
		passwordCharset = Charset.defaultCharset();
		nameToEntryMap = new HashMap<String, AesZipEntry>();

		long offsetEndOfCentralDir = this.getEndOfCentralDir();

		// TODO: This version will ignore some fields for 'disk.'

		this.sourceFile.seek(offsetEndOfCentralDir + OFFSET_NUMBER_OF_CENTRAL_DIR);
		int numOfCentralDirRecsOnDisk = this.readTwoBytesAndGetInt();
		int totalNumOfCentralDirRecs = this.readTwoBytesAndGetInt();

		if (numOfCentralDirRecsOnDisk != totalNumOfCentralDirRecs) {
			// FIXME: message: wrong form
			throw new IllegalArgumentException();
		}

		long sizeOfCentralDir = this.readFourBytesAndGetLong();
		long offsetOfStartOfCentralDir = this.readFourBytesAndGetLong();

		if (sizeOfCentralDir + offsetOfStartOfCentralDir != offsetEndOfCentralDir) {
			// FIXME: message: wrong form
			throw new IllegalArgumentException();
		}

		int sizeOfZipFileComment = this.readTwoBytesAndGetInt();
		comment = new byte[sizeOfZipFileComment];
		this.sourceFile.read(comment);

		this.sourceFile.seek(offsetOfStartOfCentralDir);
		for (int i = 0; i < totalNumOfCentralDirRecs; i++) {
			AesZipEntry entry = this.getEntryFromCentralDir();
			nameToEntryMap.put(entry.getName(), entry);
		}
		if (this.sourceFile.getFilePointer() != offsetEndOfCentralDir) {
			// FIXME: message: wrong form:
			// the end of central dirs should be the start of 'End of central
			// directory record.'
			throw new IllegalArgumentException();
		}
	}

	private AesZipEntry getEntryFromCentralDir() throws IOException {
		byte[] fourBytesBuffer = new byte[4];
		byte[] twoBytesBuffer = new byte[2];

		// central file header signature 0x02014b50 - 4 bytes
		sourceFile.read(fourBytesBuffer);
		if (!Arrays.equals(fourBytesBuffer, CENTRAL_DIR_FILE_HEADER_SIG)) {
			// FIXME: message: wrong form
			throw new IllegalArgumentException();
		}

		// ignore 'version made by (2 bytes)' and 'version needed to extract (2
		// bytes)'
		sourceFile.read(fourBytesBuffer);

		// general purpose bit flag - 2 bytes
		sourceFile.read(twoBytesBuffer);
		int generalPurposeBitFlag = ByteHelper.getInt(twoBytesBuffer);

		// compression method - 2 bytes
		sourceFile.read(twoBytesBuffer);
		int compressionMethod = ByteHelper.getInt(twoBytesBuffer);

		// last mod file time - 2 bytes
		sourceFile.read(twoBytesBuffer);
		int lastModTime = ByteHelper.getInt(twoBytesBuffer);

		// last mod file date - 2 bytes
		sourceFile.read(twoBytesBuffer);
		int lastModDate = ByteHelper.getInt(twoBytesBuffer);

		// CRC-32 - 4 bytes
		sourceFile.read(fourBytesBuffer);
		long crc = ByteHelper.getLong(fourBytesBuffer);

		// compressed size - 4 bytes
		sourceFile.read(fourBytesBuffer);
		long compressedSize = ByteHelper.getLong(fourBytesBuffer);

		// uncompressed size - 4 bytes
		sourceFile.read(fourBytesBuffer);
		long uncompressedSize = ByteHelper.getLong(fourBytesBuffer);

		// file name length - 2 bytes
		sourceFile.read(twoBytesBuffer);
		int fileNameLength = ByteHelper.getInt(twoBytesBuffer);

		// extra field length - 2 bytes
		sourceFile.read(twoBytesBuffer);
		int extraFieldLength = ByteHelper.getInt(twoBytesBuffer);

		// file comment length - 2 bytes
		sourceFile.read(twoBytesBuffer);
		int fileCommentLength = ByteHelper.getInt(twoBytesBuffer);

		// ignore 'disk number start (2 bytes)'
		sourceFile.read(twoBytesBuffer);

		// ignore 'internal file attributes (2 bytes)'
		sourceFile.read(twoBytesBuffer);

		// ignore 'external file attributes (4 bytes)'
		sourceFile.read(fourBytesBuffer);

		// relative offset of local header - 4 bytes
		sourceFile.read(fourBytesBuffer);
		long offsetOfLocalHeader = ByteHelper.getLong(fourBytesBuffer);

		// file name - variable size
		byte[] fileNameBytes = new byte[fileNameLength];
		sourceFile.read(fileNameBytes);

		// extra field - variable size
		byte[] extraField = new byte[extraFieldLength];
		sourceFile.read(extraField);

		// file comment - variable size
		byte[] fileCommentBytes = new byte[fileCommentLength];
		sourceFile.read(fileCommentBytes);

		AesZipEntry entry;
		if (AesZipDecrypter.isAesEncrypted(generalPurposeBitFlag, compressionMethod)) {
			AesZipDecrypter.checkAesExtraField(extraField);
			EncryptionStrengthMode mode = AesZipDecrypter.getEncryptionStrengthMode(extraField[8]);
			entry = new AesZipEntry(new String(fileNameBytes, this.getFileNameCharset()), mode);
			int actualMethod = ByteHelper.getInt(extraField[9], extraField[10]);
			entry.setActualCompressionMethod(actualMethod);
		} else {
			entry = new AesZipEntry(new String(fileNameBytes, this.getFileNameCharset()));
		}
		// set the values to this entry.
		entry.setCompressedSize(compressedSize);
		entry.setSize(uncompressedSize);
		entry.setMethod(compressionMethod);
		entry.setCrc(crc);
		entry.setComment(new String(fileCommentBytes, this.getCommentCharset()));
		entry.setExtra(extraField);
		entry.setRelativeOffsetOfLocalFileHeader(offsetOfLocalHeader);

		// FIXME: generate long value that represents the modification time from
		// lastModTime and lastModDate.

		return entry;
	}

	private static boolean isAesEncrypted(int generalPurposeBitFlag, int compressionMethod) {
		if ((generalPurposeBitFlag & ENCRYPTED_FLAG) == 0) {
			return false;
		}
		if (compressionMethod != CommonValues.METHOD_AES) {
			return false;
		}
		return true;
	}

	private static EncryptionStrengthMode getEncryptionStrengthMode(byte extraField) {
		switch (extraField) {
		case 1:
			return EncryptionStrengthMode.ONE;
		case 2:
			return EncryptionStrengthMode.TWO;
		case 3:
			return EncryptionStrengthMode.THREE;
		}
		// TODO
		throw new IllegalArgumentException();
	}

	private static void checkAesExtraField(byte[] extraField) {
		if (extraField.length != LENGTH_EXTRA_FIELD) {
			// FIEXME
			throw new InvalidFieldException();
		}
		for (int i = 0; i < 4; i++) {
			if (extraField[i] != AES_EXTRA_FIELD_HEADER_AND_SIZE[i]) {
				// FIEXME
				throw new InvalidFieldException();
			}
		}
		if ((extraField[4] != AES_EXTRA_FIELD_VERSION_AE1[0] && extraField[5] != AES_EXTRA_FIELD_VERSION_AE1[1])
				|| (extraField[4] != AES_EXTRA_FIELD_VERSION_AE2[0] && extraField[5] != AES_EXTRA_FIELD_VERSION_AE2[1])) {
			// FIEXME
			throw new InvalidFieldException();
		}
		if ((extraField[6] != AES_EXTRA_FIELD_VENDER_ID[0] && extraField[7] != AES_EXTRA_FIELD_VENDER_ID[1])) {
			// FIEXME
			throw new InvalidFieldException();
		}
		boolean invalidStrength = true;
		for (byte b : AES_ALLOWED_ENCRYPTION_STRENGTH) {
			if (extraField[8] == b) {
				invalidStrength = false;
				break;
			}
		}
		if (invalidStrength) {
			// FIEXME
			throw new InvalidFieldException();
		}
		int actualCompressionMethod = ByteHelper.getInt(extraField[9], extraField[10]);
		for (int method : ALLOWED_ACTUAL_METHODS) {
			if (method == actualCompressionMethod) {
				return;
			}
		}
		// FIEXME
		throw new InvalidFieldException();
	}

	private long getEndOfCentralDir() throws IOException {
		// FIXME 後ろから探して最初に見つけた 0x06054b50 を、end of central dir と仮定している。
		// 本当にそれで大丈夫かは、十分には検討していない。（他にフィールドがたまたま同じようなバイトデータで構成されていたら？
		byte[] buffer = new byte[4];
		for (long position = sourceFile.length() - 4; position > 76; position--) {
			sourceFile.seek(position);
			sourceFile.read(buffer);
			if (buffer[0] != END_OF_CENTRAL_DIR_SIG[0] || buffer[1] != END_OF_CENTRAL_DIR_SIG[1]
					|| buffer[2] != END_OF_CENTRAL_DIR_SIG[2] || buffer[3] != END_OF_CENTRAL_DIR_SIG[3]) {
				continue;
			}
			return position;
		}
		// FIXME
		throw new AesZipRuntimeException();
	}

	public Set<String> getEntryNames() {
		return Collections.unmodifiableSet(nameToEntryMap.keySet());
	}

	public AesZipEntry getEntry(String entryName) {
		return nameToEntryMap.get(entryName);
	}

	public void decryptAndExtractAllFiles(File destinationBaseDirectory, String password) throws IOException,
			PasswordNotCorrectException {
		String[] entryNames = getEntryNames().toArray(new String[nameToEntryMap.size()]);
		this.decryptAndExtractFiles(entryNames, destinationBaseDirectory, password);
	}

	public void decryptAndExtractFiles(String[] entryNames, File destinationBaseDirectory, String password)
			throws IOException, PasswordNotCorrectException {
		if (password == null || destinationBaseDirectory == null) {
			// TODO
			throw new IllegalArgumentException(new NullPointerException());
		}
		if (!destinationBaseDirectory.isDirectory()) {
			// TODO
			throw new IllegalArgumentException();
		}
		if (progressListener != null) {
			progressListener.start();
			long totalSize = 0;
			for (String entryName : entryNames) {
				totalSize += nameToEntryMap.get(entryName).getContentCompressedSize();
			}
			previousStatus = new ArchiveProgressStatus(entryNames.length, 0, totalSize, 0, -1, 0);
		}
		for (String entryName : entryNames) {
			AesZipEntry entry = nameToEntryMap.get(entryName);
			File destination = new File(destinationBaseDirectory, entryName);
			if (progressListener != null) {
				long size = entry.getContentCompressedSize();
				previousStatus = previousStatus.generateUpdatedStatusWithNextEntry(size, entryName,
						destination.getCanonicalPath());
				progressListener.notifyStatus(previousStatus);
			}
			if (entry.isDirectory()) {
				// if the entry is a directory, it only generates the directory.
				// TODO: check entry contents?
				destination.mkdirs();
				continue;
			}

			this.prepare(entry, password);
			// prepare the stream
			if (!destinationBaseDirectory.exists()) {
				destinationBaseDirectory.mkdirs();
			}
			destination.getParentFile().mkdirs();
			OutputStream destinationStream;
			try {
				destinationStream = new FileOutputStream(destination);
			} catch (FileNotFoundException e) {
				// TODO: more comprehensive message.
				throw new IllegalArgumentException(e);
			}
			try {
				this.decryptAndExtractImpl(entry, destinationStream);
			} finally {
				if (destinationStream != null) {
					destinationStream.close();
				}
			}
		}
		if (progressListener != null) {
			progressListener.finishedAll();
		}
	}

	/**
	 * Decrypt and extract a file. If the entry is a directory, this method will simply return (do nothing). If
	 * the <code>destination</code> does not exist, this method will create the
	 * directory or the file.
	 * 
	 * @param entryName
	 * @param destination
	 * @param password
	 * @throws IOException
	 * @throws PasswordNotCorrectException
	 */
	public void decryptAndExtract(String entryName, File destination, String password) throws IOException,
			PasswordNotCorrectException {
		// check arguments
		this.checkEntryName(entryName);
		if (password == null || destination == null) {
			// TODO
			throw new IllegalArgumentException(new NullPointerException());
		}
		if (destination.isDirectory()) {
			// TODO
			throw new IllegalArgumentException();
		}

		AesZipEntry entry = nameToEntryMap.get(entryName);
		if (progressListener != null) {
			progressListener.start();
			long size = entry.getContentCompressedSize();
			previousStatus = new ArchiveProgressStatus(1, 0, size, 0, -1, 0);
		}
		if (entry.isDirectory()) {
			return;
		}

		this.prepare(entry, password);

		// prepare the stream
		File destinationBaseDirectory = destination.getCanonicalFile().getParentFile();
		if (!destinationBaseDirectory.exists()) {
			destinationBaseDirectory.mkdirs();
		}
		OutputStream destinationStream;
		try {
			destinationStream = new FileOutputStream(destination);
		} catch (FileNotFoundException e) {
			// TODO: more comprehensive message.
			throw new IllegalArgumentException(e);
		}
		try {
			this.decryptAndExtractImpl(entry, destinationStream);
		} finally {
			if (destinationStream != null) {
				destinationStream.close();
			}
		}
		if (progressListener != null) {
			progressListener.finishedAll();
		}
	}

	public void decryptAndExtract(String entryName, OutputStream destinationStream, String password)
			throws IOException, PasswordNotCorrectException {
		// check arguments
		this.checkEntryName(entryName);
		if (password == null || destinationStream == null) {
			// TODO
			throw new IllegalArgumentException(new NullPointerException());
		}

		AesZipEntry entry = nameToEntryMap.get(entryName);
		if (progressListener != null) {
			progressListener.start();
			long size = entry.getContentCompressedSize();
			previousStatus = new ArchiveProgressStatus(1, 0, size, 0, -1, 0);
		}
		if (entry.isDirectory()) {
			// no op. if the entry is a directory.
			return;
		}

		this.prepare(entry, password);
		this.decryptAndExtractImpl(entry, destinationStream);
		if (progressListener != null) {
			progressListener.finishedAll();
		}
	}

	private void decryptAndExtractImpl(AesZipEntry entry, OutputStream destinationStream) throws IOException {
		OutputStream targetStream;
		short actualMethod = entry.getActualCompressionMethod();
		switch (actualMethod) {
		case CommonValues.METHOD_DEFLATED:
			targetStream = new InflaterOutputStream(destinationStream, new Inflater(true));
			break;
		case CommonValues.METHOD_STORED:
			targetStream = destinationStream;
			break;
		default:
			throw new InvalidMethodException(actualMethod);
		}

		AesCtrBlockCipherOutputStream cipherOut = AesCtrBlockCipherOutputStream.generateDecryptInstance(targetStream);
		MacFilterOutputStream macFilter = new MacFilterOutputStream(cipherOut);
		cipherOut.init(entry.getEncryptionKey());
		macFilter.init(entry.getAuthenticationKey());

		// decrypt and uncompress
		long size = entry.getContentCompressedSize();

		int readSize = BLOCK_SIZE;

		for (; readSize < size; readSize += BLOCK_SIZE) {
			sourceFile.read(decryptionBlock, 0, BLOCK_SIZE);
			macFilter.write(decryptionBlock);
			updateProgress(decryptionBlock.length);
		}
		int rest = (int) (size - readSize) + BLOCK_SIZE;
		if (rest > 1) {
			sourceFile.read(decryptionBlock, 0, rest);
			macFilter.write(decryptionBlock, 0, rest);
			updateProgress(rest);
		}
		byte[] authCode = new byte[LENGTH_AUTHENTICATION_CODE];
		sourceFile.read(authCode, 0, LENGTH_AUTHENTICATION_CODE);

		macFilter.flush();
		cipherOut.flush();
		targetStream.flush();
		destinationStream.flush();

		// check. FIXME
		// AE-2 の場合は crcチェックはしない。authcodeはAE-1もAE-2もチェックする。当面は AE-2のみ。
		// long crc = macFilter.getCrc();
		// サイズもチェックする？
		byte[] authCodeExpected = macFilter.getAuthenticationCode();
		if (!Arrays.equals(authCode, authCodeExpected)) {
			// FIXME
			throw new AesZipRuntimeException(MessageFormat.format("MAC value not match for [{0}]", entry.getName()));
		}
	}

	private void prepare(AesZipEntry entry, String password) throws IOException, PasswordNotCorrectException {
		sourceFile.seek(entry.getRelativeOffsetOfLocalFileHeader() + LENGTH_LOCAL_FILE_HEADER_SKIP);

		// FIXME: compare with central dir(?)
		int fileNameLength = this.readTwoBytesAndGetInt();
		int extraFieldLength = this.readTwoBytesAndGetInt();
		sourceFile.skipBytes(fileNameLength + extraFieldLength);

		EncryptionStrengthMode mode = entry.getStrengthMode();
		int lengthSalt = mode.getSaltLength();
		byte[] saltValue = new byte[lengthSalt];
		byte[] verificationValue = new byte[LENGTH_PASSWORD_VERIFICATION_VALUE];
		sourceFile.read(saltValue, 0, lengthSalt);
		sourceFile.read(verificationValue, 0, LENGTH_PASSWORD_VERIFICATION_VALUE);
		entry.deriveKeys(password.getBytes(this.getPasswordCharset()), saltValue, verificationValue);
	}

	public void setPasswordCharset(Charset charset) {
		if (charset == null) {
			// TODO: more comprehensive message.
			throw new IllegalArgumentException(new NullPointerException());
		}
		passwordCharset = charset;
	}

	private Charset getPasswordCharset() {
		return passwordCharset;
	}

	public void setFileNameCharset(Charset charset) {
		if (charset == null) {
			// TODO: more comprehensive message.
			throw new IllegalArgumentException(new NullPointerException());
		}
		fileNameCharset = charset;
	}

	private Charset getFileNameCharset() {
		return fileNameCharset;
	}

	public void setCommentCharset(Charset charset) {
		if (charset == null) {
			// TODO: more comprehensive message.
			throw new IllegalArgumentException(new NullPointerException());
		}
		commentCharset = charset;
	}

	private Charset getCommentCharset() {
		return commentCharset;
	}

	private void checkEntryName(String entryName) {
		if (entryName == null || !nameToEntryMap.containsKey(entryName)) {
			// TODO
			throw new IllegalArgumentException(new NullPointerException());
		}
	}

	public void setProgressListener(ArchiveProgressListener progressListener) {
		this.progressListener = progressListener;
	}

	private void updateProgress(long increments) {
		if (progressListener != null) {
			previousStatus = previousStatus.generateUpdatedStatus(increments);
			progressListener.notifyStatus(previousStatus);
		}
	}

	private int readTwoBytesAndGetInt() throws IOException {
		byte[] buffer = new byte[2];
		this.sourceFile.read(buffer);
		return ByteHelper.getInt(buffer);
	}

	private long readFourBytesAndGetLong() throws IOException {
		byte[] buffer = new byte[4];
		this.sourceFile.read(buffer);
		return ByteHelper.getLong(buffer);
	}
}
