package net.y3n20u.aeszip;

import static net.y3n20u.aeszip.CommonValues.CENTRAL_DIR_START_DISK_NUMBER;
import static net.y3n20u.aeszip.CommonValues.CENTRAL_FILE_HEADER_SIG;
import static net.y3n20u.aeszip.CommonValues.DATA_DESCRIPTOR_SIG;
import static net.y3n20u.aeszip.CommonValues.DEFLATER_LEVEL;
import static net.y3n20u.aeszip.CommonValues.END_OF_CENTRAL_SIG;
import static net.y3n20u.aeszip.CommonValues.EXTERNAL_FILE_ATTRIBUTE;
import static net.y3n20u.aeszip.CommonValues.FILE_NAME_SEPARATOR;
import static net.y3n20u.aeszip.CommonValues.FILE_START_DISK_NUMBER;
import static net.y3n20u.aeszip.CommonValues.FLAG_FOR_DIRECTORY;
import static net.y3n20u.aeszip.CommonValues.FLAG_FOR_ENCRYPTED_FILE;
import static net.y3n20u.aeszip.CommonValues.INTERNAL_FILE_ATTRIBUTE;
import static net.y3n20u.aeszip.CommonValues.LOCAL_FILE_HEADER_SIG;
import static net.y3n20u.aeszip.CommonValues.MAX_TWO_BYTE_FIELD;
import static net.y3n20u.aeszip.CommonValues.MESSAGE_ALREADY_EXISTS;
import static net.y3n20u.aeszip.CommonValues.MESSAGE_COMMENT_LENGTH_TOOLONG;
import static net.y3n20u.aeszip.CommonValues.MESSAGE_NAMELENGTH_TOOLONG;
import static net.y3n20u.aeszip.CommonValues.MESSAGE_RECORDS_TOO_MUCH;
import static net.y3n20u.aeszip.CommonValues.MESSAGE_WRITE_WHILE_NO_ENTRY;
import static net.y3n20u.aeszip.CommonValues.METHOD_AES;
import static net.y3n20u.aeszip.CommonValues.METHOD_DEFLATED;
import static net.y3n20u.aeszip.CommonValues.METHOD_STORED;
import static net.y3n20u.aeszip.CommonValues.NUMBER_OF_THIS_DISK;
import static net.y3n20u.aeszip.CommonValues.VERSION_DEFLATED;
import static net.y3n20u.aeszip.CommonValues.VERSION_STORED;

import java.io.FilterOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.charset.Charset;
import java.text.MessageFormat;
import java.util.ArrayList;
import java.util.List;
import java.util.zip.Deflater;
import java.util.zip.DeflaterOutputStream;

/**
 * @see <a href="http://www.gladman.me.uk/cryptography_technology/fileencrypt/">BRG Main SIte</a>
 * @see <a href="http://www.winzip.com/aes_info.htm">WinZip® - AES Encryption Information</a>
 * 
 * @author y3n20u@gmail.com
 */
public class AesZipOutputStream extends FilterOutputStream {
	// TODO: AE-1.
	// TODO: general purpose flag の第3ビットが立たない方式（事前にサイズやCRCがわかる場合）
	// TODO: speed up.
	// TODO: file name, comment encoding. entryに持たせる？
	
	private Deflater deflater;
	private DeflaterOutputStream deflaterOut;
	private final AesCtrBlockCipherOutputStream cipherOut;
	private final MacFilterOutputStream macFilter;
	private final List<AesZipEntry> entryList;
	private AesZipEntry currentEntry;
	private int totalWrittenSize = 0;
	private String comment;
	private OutputStream targetStream;
	private Charset fileNameCharset = Charset.defaultCharset();
	private Charset commentCharset = Charset.defaultCharset();
	private boolean finished = false;

	public AesZipOutputStream(OutputStream out) {
		super(out);
		macFilter = new MacFilterOutputStream(this.out);
		cipherOut = AesCtrBlockCipherOutputStream.generateEncryptInstance(macFilter);
		entryList = new ArrayList<AesZipEntry>();
	}

	@Override
	public synchronized void write(byte[] b) throws IOException {
		this.write(b, 0, b.length);
	}

	@Override
	public synchronized void write(byte[] b, int off, int len) throws IOException {
		if (currentEntry == null) {
			throw new IllegalStateException(MESSAGE_WRITE_WHILE_NO_ENTRY);
		}
		targetStream.write(b, off, len);
		long writtenSize = currentEntry.getSize() + len;
		currentEntry.setSize(writtenSize);
	}

	/**
	 * Set the comment for this zip file. The length of the comment must be
	 * smaller than 65536 (because the field for comment-length is two-byte).
	 * 
	 * @param comment
	 *            comment for this zip file.
	 * @throws IllegalArgumentException
	 *             the comment is too long.
	 */
	public void setZipFileComment(String comment) {
		if (comment == null) {
			comment = "";
		} else {
			int commentLength = comment.getBytes(this.getCommentCharset()).length;
			if (commentLength > MAX_TWO_BYTE_FIELD) {
				throw new IllegalArgumentException(MessageFormat.format(MESSAGE_COMMENT_LENGTH_TOOLONG, commentLength));
			}
		}
		this.comment = comment;
	}

	public void putNextZipEntry(AesZipEntry nextEntry) throws IOException {
		if (nextEntry == null) {
			// TODO: more comprehensive message.
			throw new IllegalArgumentException(new NullPointerException());
		}
		int method = nextEntry.getMethod();
		if (method != METHOD_AES && method != METHOD_STORED) {
			throw new InvalidMethodException(method);
		}

		if (this.hasSameNameEntry(nextEntry.getName())) {
			throw new IllegalArgumentException(MessageFormat.format(MESSAGE_ALREADY_EXISTS, nextEntry.getName()));
		}

		if (currentEntry != null) {
			closeZipEntry();
		}

		nextEntry.setRelativeOffsetOfLocalFileHeader(totalWrittenSize);
		if (nextEntry.isDirectory()) {
			// if the entry indicates directory,
			// the method should be 'STORED' and the entry should not to be
			// encrypted.
			nextEntry.setMethod(METHOD_STORED);
			nextEntry.setCompressedSize(0);
			nextEntry.setSize(0);
			nextEntry.setCrc(0);
			entryList.add(nextEntry);
			writeLocalFileHeader(nextEntry);
			currentEntry = null;
		} else { // TODO: impl. case with no encryption
			writeLocalFileHeader(nextEntry);
			// write the salt value.
			this.writeByteArray(nextEntry.getSaltValue());
			// write the password verification value
			this.writeByteArray(nextEntry.getPasswordVerificationValue());
			nextEntry.setSize(0L);
			currentEntry = nextEntry;
			this.initStream();
		}
	}

	private void initStream() {
		macFilter.init(currentEntry.getAuthenticationKey());
		cipherOut.init(currentEntry.getEncryptionKey());
		short actualMethod = currentEntry.getActualCompressionMethod();
		switch (actualMethod) {
		case CommonValues.METHOD_DEFLATED:
			deflater = new Deflater(DEFLATER_LEVEL, true);
			deflaterOut = new DeflaterOutputStream(cipherOut, deflater);
			targetStream = deflaterOut;
			break;
		case CommonValues.METHOD_STORED:
			targetStream = cipherOut;
			break;
		default:
			throw new InvalidMethodException(actualMethod);
		}
	}

	@Override
	public void close() throws IOException {
		try {
			if (!finished) {
				this.finish();
			}
		}
		finally {
			super.close();
		}
	}

	public void finish() throws IOException {
		finished = true;
		
		// 1. close the previous entry.
		if (currentEntry != null) {
			this.closeZipEntry();
		}
		
		// 2. write the central directories.
		int startOfCentralDir = totalWrittenSize;
		for (AesZipEntry aesEntry : entryList) {
			this.writeCentralFileHeader(aesEntry);
		}
		int sizeOfCentralDir = totalWrittenSize - startOfCentralDir;
		this.writeEndOfCentralFileHeader(sizeOfCentralDir, startOfCentralDir);
	}

	public void closeZipEntry() throws IOException {
		if (currentEntry == null) {
			// no operation.
			return;
		}
		entryList.add(currentEntry);

		long compressedSize;
		short actualMethod = currentEntry.getActualCompressionMethod();
		switch (actualMethod) {
		case METHOD_DEFLATED:
			deflaterOut.finish();
			compressedSize = deflater.getBytesWritten();
			break;
		case METHOD_STORED:
			compressedSize = currentEntry.getSize();
			break;
		default:
			throw new InvalidMethodException(actualMethod);
		}
		totalWrittenSize += compressedSize;
		cipherOut.flush();

		currentEntry.setContentCompressedSize(compressedSize);
		currentEntry.setCrc(macFilter.getCrc());

		// write the authentication code.
		this.writeByteArray(macFilter.getAuthenticationCode());

		// general purpose flag の第3ビットが立っているなら、ここでサイズやCRCを書く。
		if ((AesZipOutputStream.generateFlagValue(currentEntry) & 0x08) == 0x08) {
			this.writeDataDescriptor(currentEntry);
		}
		currentEntry = null;
	}

	private void writeDataDescriptor(AesZipEntry entry) throws IOException {
		// signature - 4 bytes
		this.writeLongInLittleEndian(DATA_DESCRIPTOR_SIG, 4);

		// CRC-32 - 4 bytes
		// TODO: if (AE-1?) { this.writeLongInLittleEndian(entry.getCrc(), 4);
		// }
		this.writeLongInLittleEndian(0, 4);

		// compressed size - 4 bytes
		this.writeLongInLittleEndian(entry.getCompressedSize(), 4);

		// uncompressed size - 4 bytes
		this.writeLongInLittleEndian(entry.getSize(), 4);
	}

	private void writeLocalFileHeader(AesZipEntry entry) throws IOException {
		// local file header signature = 0x04034b50 - 4 bytes
		this.writeLongInLittleEndian(LOCAL_FILE_HEADER_SIG, 4);

		this.writeCommonFileHeader1(entry);
		// CRC - 4 bytes
		// compressed size. - 4 bytes
		// uncompressed size. - 4 bytes
		if (((AesZipOutputStream.generateFlagValue(entry) & 0x08) == 0x08) || entry.isDirectory()) {
			this.write4ByteFieldInLittleEndian(0);
			this.write4ByteFieldInLittleEndian(0);
			this.write4ByteFieldInLittleEndian(0);
		} else {
			this.write4ByteFieldInLittleEndian(0); // TODO: for AE-1?
			this.write4ByteFieldInLittleEndian((int) entry.getCompressedSize());
			this.write4ByteFieldInLittleEndian((int) entry.getSize());
		}

		this.writeCommonFileHeader2(entry);
		this.writeCommonFileHeader3(entry);
	}

	private void writeCentralFileHeader(AesZipEntry aesEntry) throws IOException {
		// central file header signature = 0x02014b50 - 4 bytes
		this.writeLongInLittleEndian(CENTRAL_FILE_HEADER_SIG, 4);

		// version made by. - 2 bytes
		this.write2ByteFieldInLittleEndian(AesZipOutputStream.getVersion(aesEntry));

		this.writeCommonFileHeader1(aesEntry);

		// CRC - 4 bytes
		this.write4ByteFieldInLittleEndian(0); // TODO: for AE-1?

		// compressed size. - 4 bytes
		// uncompressed size. - 4 bytes
		this.write4ByteFieldInLittleEndian((int) aesEntry.getCompressedSize());
		this.write4ByteFieldInLittleEndian((int) aesEntry.getSize());

		this.writeCommonFileHeader2(aesEntry);

		// file comment length - 2 bytes
		String comment = aesEntry.getComment();
		int commentLength = ((comment == null) ? 0 : comment.getBytes(this.getCommentCharset()).length);
		this.write2ByteFieldInLittleEndian((short) commentLength);

		// disk number where file starts - 2 bytes
		this.write2ByteFieldInLittleEndian(FILE_START_DISK_NUMBER);

		// internal file attributes. - 2 bytes
		this.write2ByteFieldInLittleEndian(INTERNAL_FILE_ATTRIBUTE);

		// external file attributes. - 4 bytes
		this.write4ByteFieldInLittleEndian(EXTERNAL_FILE_ATTRIBUTE);

		// relative offset of local file header - 4 bytes
		this.writeLongInLittleEndian(aesEntry.getRelativeOffsetOfLocalFileHeader(), 4);

		this.writeCommonFileHeader3(aesEntry);

		// file comment - variable
		this.writeByteArray((comment != null) ? comment.getBytes(this.getCommentCharset()) : new byte[0]);
	}

	private void writeCommonFileHeader1(AesZipEntry entry) throws IOException {
		// version needed to extract (minimum). - 2 bytes
		this.write2ByteFieldInLittleEndian(AesZipOutputStream.getVersion(entry));

		// general purpose flag - 2 bytes
		// 暗号化しているので第0ビットを立てる。サイズ不明の場合は第3ビットも立てる。
		this.write2ByteFieldInLittleEndian(AesZipOutputStream.generateFlagValue(entry));

		// compression method: 99 (AES, 0x63) or 0 (not encrypted) - 2 bytes
		this.write2ByteFieldInLittleEndian((short) entry.getMethod());

		// file last modification time - 2 bytes
		this.write2ByteFieldInLittleEndian(entry.getLastModTime());

		// file last modification date - 2 bytes
		this.write2ByteFieldInLittleEndian(entry.getLastModDate());

	}

	private void writeCommonFileHeader2(AesZipEntry entry) throws IOException {
		// file name length - 2 bytes
		byte[] nameByteArray = entry.getName().getBytes(this.getFileNameCharset());
		if (nameByteArray.length > MAX_TWO_BYTE_FIELD) {
			throw new IllegalArgumentException(MessageFormat.format(MESSAGE_NAMELENGTH_TOOLONG, nameByteArray.length));
		}
		this.write2ByteFieldInLittleEndian((short) nameByteArray.length);

		// extra field length - 2 bytes
		byte[] extra = entry.getExtra();
		short extraFieldLength = (short) ((extra != null ? extra.length : 0));
		this.write2ByteFieldInLittleEndian(extraFieldLength);
	}

	private void writeCommonFileHeader3(AesZipEntry entry) throws IOException {
		// file name - variable
		this.writeByteArray(entry.getName().getBytes(this.getFileNameCharset()));

		// extra field - variable
		byte[] extra = entry.getExtra();
		if (extra != null) {
			this.out.write(extra);
			totalWrittenSize += extra.length;
		}
	}

	/**
	 * generate the value of "general purpose bit flag."
	 * <p>
	 * TODO: not understand details of the spec.
	 * 
	 * @param entry
	 * @return
	 */
	private static short generateFlagValue(AesZipEntry entry) {
		if (entry == null) {
			throw new IllegalArgumentException(new NullPointerException());
		}
		return entry.isDirectory() ? FLAG_FOR_DIRECTORY : FLAG_FOR_ENCRYPTED_FILE;
	}

	private void writeEndOfCentralFileHeader(int sizeOfCentralDir, int startOfCentralDir) throws IOException {
		// 3 write the end of the central directory record
		// end of central directory signature = 0x06054b50 - 4 bytes
		this.writeLongInLittleEndian(END_OF_CENTRAL_SIG, 4);

		// Number of this disk - 2 bytes
		this.write2ByteFieldInLittleEndian(NUMBER_OF_THIS_DISK);

		// disk where central directory starts - 2 bytes
		this.write2ByteFieldInLittleEndian(CENTRAL_DIR_START_DISK_NUMBER);

		// number of central directory records on this disk - 2 bytes
		this.write2ByteFieldInLittleEndian(getEntryListSize());

		// total number of central directory records - 2 bytes
		this.write2ByteFieldInLittleEndian(getEntryListSize());

		// size of central directory (bytes) - 4 bytes
		this.write4ByteFieldInLittleEndian(sizeOfCentralDir);

		// offset of start of central directory, relative to start of archive -
		// 4 bytes
		this.write4ByteFieldInLittleEndian(startOfCentralDir);

		// zip file comment length - 2 bytes
		// the length must have been already checked (won't be checked here).
		byte[] commentByteArray = ((comment == null) ? new byte[0] : comment.getBytes(this.getCommentCharset()));
		this.write2ByteFieldInLittleEndian((short) commentByteArray.length);

		// zip file comment - variable
		this.writeByteArray(commentByteArray);
	}

	private short getEntryListSize() {
		int numberOfRecords = entryList.size();
		if (numberOfRecords > MAX_TWO_BYTE_FIELD || numberOfRecords < 0) {
			throw new IllegalStateException(MessageFormat.format(MESSAGE_RECORDS_TOO_MUCH, numberOfRecords));
		}
		return (short) numberOfRecords;
	}

	private boolean hasSameNameEntry(String name) {
		String comparedName = name;
		if (!name.endsWith(FILE_NAME_SEPARATOR)) {
			comparedName = name + FILE_NAME_SEPARATOR;
		}
		for (AesZipEntry entry : entryList) {
			if (entry.getName().endsWith(FILE_NAME_SEPARATOR)) {
				if (comparedName.equalsIgnoreCase(entry.getName())) {
					return true;
				}
			} else {
				if (comparedName.equalsIgnoreCase(entry.getName() + FILE_NAME_SEPARATOR)) {
					return true;
				}
			}
		}
		return false;
	}
	
	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());
		}
		this.commentCharset = charset;
	}

	private Charset getCommentCharset() {
		return commentCharset;
	}

	private static short getVersion(AesZipEntry entry) {
		short actualMethod = entry.getActualCompressionMethod();
		switch (actualMethod) {
		case METHOD_DEFLATED:
			return VERSION_DEFLATED;
		case METHOD_STORED:
			return VERSION_STORED;
		default:
			throw new InvalidMethodException(actualMethod);
		}
	}

	private void writeByteArray(byte... values) throws IOException {
		this.out.write(values);
		totalWrittenSize += values.length;
	}

	private void writeLongInLittleEndian(long value, int count) throws IOException {
		long v = value;
		for (int i = 0; i < count; i++) {
			this.out.write((byte) (v & 0xff));
			v >>>= 8;
		}
		totalWrittenSize += count;
	}

	private void write2ByteFieldInLittleEndian(short value) throws IOException {
		this.out.write((byte) (value & 0xff));
		this.out.write((byte) ((value >>> 8) & 0xff));
		totalWrittenSize += 2;
	}

	private void write4ByteFieldInLittleEndian(int value) throws IOException {
		this.out.write((byte) (value & 0xff));
		this.out.write((byte) ((value >>> 8) & 0xff));
		this.out.write((byte) ((value >>> 16) & 0xff));
		this.out.write((byte) ((value >>> 24) & 0xff));
		totalWrittenSize += 4;
	}
}
