package org.expeditee.encryption.io;

import java.io.IOException;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.util.Arrays;
import java.util.Base64;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.stream.Collectors;

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

import org.expeditee.encryption.CryptographyConstants;
import org.expeditee.encryption.items.surrogates.EncryptionDetail;
import org.expeditee.encryption.items.surrogates.EncryptionDetail.Type;
import org.expeditee.encryption.items.surrogates.Label.LabelResult;
import org.expeditee.encryption.items.surrogates.Label;
import org.expeditee.encryption.items.surrogates.Label.LabelInfo;
import org.expeditee.gui.Frame;
import org.expeditee.gui.Frame.BodyType;
import org.expeditee.gui.MessageBay;
import org.expeditee.io.Conversion;
import org.expeditee.io.ExpWriter;
import org.expeditee.items.Item;
import org.expeditee.items.Line;
import org.expeditee.items.Text;
import org.expeditee.settings.identity.secrets.KeyList;
import org.expeditee.stats.SessionStats;

public class EncryptedExpWriter extends ExpWriter implements CryptographyConstants {
	private SecretKey key;
	
	private final String label;
	
	private static final String ENCRYPTED_EXP_FLAG = "EncryptedExp";
	private static final String nl = "\n";
	private static final String labelProfile = "Profile";
	private static final String labelNone = "None";
	protected static final char TERMINATOR = 'Z';
	protected static final String TERMINATOR_WITH_CONTINUATION = TERMINATOR + "...";

	public EncryptedExpWriter(String encryptionLabel) {
		label = encryptionLabel;
		if (label.equals(getLabelNone())) {
			return;
		}
		
		if (label.equals(labelProfile)) {
			// obtain personal key
			Text text = KeyList.PersonalKey.get();
			List<String> data = text.getData();
			if (data != null && !data.isEmpty()) {
				byte[] keyBytes = Base64.getDecoder().decode(data.get(0));
				key = new SecretKeySpec(keyBytes, SymmetricAlgorithm);
			}
		} else {
			LabelInfo res = Label.getLabel(label);
			if (res.is(LabelResult.SuccessResolveLabelToKey)) {
				byte[] keyBytes = res.key;
				key = new SecretKeySpec(keyBytes, SymmetricAlgorithm);
			} else if (res.is(LabelResult.ErrorUnableToFindLabel)) {
				MessageBay.errorMessage(res.toString() + encryptionLabel);
				key = null;
			} else {
				MessageBay.errorMessage(res.toString());
			}
		}
	}
	
	@Override
	protected void preOutputFrame() {
		try {
			String line = ENCRYPTED_EXP_FLAG + label + nl;
			_writer.write(line);
			_stringWriter.append(line);
		} catch (final IOException e) { 
			e.printStackTrace();
		}
	}
	
	@Override
	public void outputFrame(Frame frame) throws IOException {
		if (_writer == null) { return; }
		
		preOutputFrame();
		writeHeader(frame);
		
		// write item
		writeItemData(frame, true);
		writeTerminator();
		// write lines and constraints
		writeLineData();
		writeTerminator();
		writeConstraintData();
		
		// write surrogate items
		if (frame.hasSurrogates()) {
			writeSurrogateTerminator();
			writeItemData(frame, false);
			writeTerminator();
		} else {
			writeTerminator();
		}
				
		writeLine(SessionStats.getFrameEventList(frame));
	}
	
	@Override
	protected void writeLine(String line) throws IOException {
		// do not write empty lines
		if (line == null) { return; }
		
		String toWrite;
		if (key == null && label.equals(getLabelNone())) {
			toWrite = line + nl;
		} else {
			// prepare line to write out
			byte[] encrypted = EncryptSymmetric(line.getBytes(), key);
			toWrite = Base64.getEncoder().encodeToString(encrypted) + nl;
		}
				
		// output
		_writer.write(toWrite);
		_stringWriter.append(toWrite);
	}
	
	protected void writeClass(Item toWrite) throws IOException {
		LinkedHashMap<Character,Method> itemTags = new LinkedHashMap<Character, Method>(getItemCharTags());
		LinkedHashMap<String,Method> itemTagsExt = new LinkedHashMap<String, Method>(getItemStrTags());
		
		writeTag(toWrite, new Object[] {}, itemTags, 'S');
		writeTag(toWrite, new Object[] {}, itemTagsExt, "_el");
		
		if (toWrite.isSurrogate()) {
			writeLine("SurrogateFor " + toWrite.getPrimary().getID());
		}
		
		itemTags.remove('S');
		itemTagsExt.remove("_el");
				
		writeTags(toWrite, new Object[] {}, itemTags);
		writeTags(toWrite, new Object[] {}, itemTagsExt);
	}
	
	@Override
	protected <T> void writeTag(Item toWrite, Object[] param, LinkedHashMap<T, Method> tags, T tag) {
		if (toWrite.isSurrogate() && toWrite.isTagInherited(tag + "")) {
			return;
		}
		EncryptionDetail encryptionDetail = toWrite.getEncryptionDetailForTag(tag + "");
		Type encryptionDetailType = encryptionDetail.getEncryptionDetailType();
				
		switch(encryptionDetailType) {
		case UnencryptedOnSave:
			writeTagUnencryptedOnSave(toWrite, param, tags, tag);
			break;
		case InheritanceCheckOnSave:
			writeTagInheritanceCheckOnSave(toWrite, tags, tag);
			break;
		case ReencryptOnSave:
			writeTagReencryptOnSave(toWrite, tags, tag);
			break;
		case UseUndecipheredValueOnSave:
			writeTagUseUndecipheredValueOnSave(toWrite, tags, tag);
			break;
		}
	}
	
	private void writeItemData(Frame frame, boolean primaryItems) throws IOException {
		// write each item in the frame
		BodyType bodyType = primaryItems ? BodyType.PrimaryBody : BodyType.SurrogateBody;
		for (Item i : frame.getItemsToSave(bodyType)) {
			assert (!(i instanceof Line));
			writeItem(i);
		}
		
		for (Item i: frame.getBodyItemsWithInsufficientPermissions()) {
			assert (!(i instanceof Line));
			writeItem(i);
		}
	}
	
	private <T> void writeTagInheritanceCheckOnSave(Item toWrite, LinkedHashMap<T, Method> tags, T tag) {	
		List<Item> surrogateItems = toWrite.getSurrogates().stream().collect(Collectors.toList());
		Function<Item, Boolean> isTagInherited = surrogate -> surrogate.isTagInherited(tag + "");
		BinaryOperator<Boolean> trueExists = (a, b) -> a || b;
		boolean surrogatesInherit = surrogateItems.stream().map(isTagInherited).collect(Collectors.reducing(trueExists)).orElseGet(() -> false);
		boolean userHasKey = Label.getLabel(toWrite.getEncryptionLabel()).is(LabelResult.SuccessResolveLabelToKey);
		
		// If we have no surrogates that inherit this property from us, and we have the label required to encrypt it, then we should encrypt it.
		if (!surrogatesInherit && userHasKey) {
			EncryptionDetail reencryptOnSave = new EncryptionDetail(EncryptionDetail.Type.ReencryptOnSave);
			toWrite.setEncryptionDetailForTag(tag + "", reencryptOnSave);
			writeTagReencryptOnSave(toWrite, tags, tag);
		} else {
			// If one or more surrogates still inherit from us, or we do not have the label, then we write it out unencrypted.
			EncryptionDetail unencryptedOnSave = new EncryptionDetail(EncryptionDetail.Type.UnencryptedOnSave);
			toWrite.setEncryptionDetailForTag(tag + "", unencryptedOnSave);
			writeTagUnencryptedOnSave(toWrite, new Object[] {}, tags, tag);
		}
	}
	
	private <T> void writeTagReencryptOnSave(Item toWrite, LinkedHashMap<T, Method> tags, T tag) {
		Method toRun = tags.get(tag);
		Class<?> declarer = toRun.getDeclaringClass();
		LabelInfo res = Label.getLabel(toWrite.getEncryptionLabel());
		if (declarer.isAssignableFrom(toWrite.getClass()) && res.is(LabelResult.SuccessResolveLabelToKey)) {
			try {
				Object o = toRun.invoke(toWrite, new Object[] {});
				o = Conversion.ConvertToExpeditee(toRun, o);
				if (o == null) {
					return;
				}
				if (o instanceof List) {
					for (Object line: (List<?>) o) {
						writeLineEnc(tag, line, res.key);
					}
				} else {
					writeLineEnc(tag, o, res.key);
				}
			} catch (IllegalAccessException | IllegalArgumentException | InvocationTargetException e) {
				e.printStackTrace();
			} catch (IOException e) {
				e.printStackTrace();
			}
		}
	}
	
	private <T> void writeTagUseUndecipheredValueOnSave(Item toWrite, LinkedHashMap<T, Method> tags, T tag) {
		EncryptionDetail encryptionDetail = toWrite.getEncryptionDetailForTag(tag + "");
		try {
			writeLine(tag + "E", encryptionDetail.getUndecipheredValue());
		} catch (IOException e) {
			e.printStackTrace();
		}
	}
	
	private <T> void writeTagUnencryptedOnSave(Item toWrite, Object[] param, LinkedHashMap<T, Method> tags, T tag) {
		super.writeTag(toWrite, param, tags, tag);
	}
	
	private void writeSurrogateTerminator() throws IOException {
		writeLine(TERMINATOR_WITH_CONTINUATION + nl);
	}
	
	private <T> void writeLineEnc(T tag, Object line, byte[] keyBytes) throws IOException {
		SecretKey key = new SecretKeySpec(keyBytes, SymmetricAlgorithm);
		byte[] lineEncryptedBytes = EncryptSymmetric(line.toString().getBytes(), key);
		line = Base64.getEncoder().encodeToString(lineEncryptedBytes);
		writeLine(tag + "E", line.toString());
	}
	
	private static byte[] EncryptSymmetric(byte[] toEncrypt, SecretKey key) {
		try {
			Cipher cipher = Cipher.getInstance(SymmetricAlgorithm + SymmetricAlgorithmParameters);
			cipher.init(Cipher.ENCRYPT_MODE, key);
			//could use modulus
			int length = (int) ((Math.ceil(toEncrypt.length / 16f)) * 16);
			byte[] toEncryptSizeAdjusted = Arrays.copyOf(toEncrypt, length);
			byte[] result = cipher.doFinal(toEncryptSizeAdjusted);
			return result;
		} catch (final NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException
				| IllegalBlockSizeException | BadPaddingException e) {
			e.printStackTrace();
			return null;
		}
	}

	/**
	 * @return the labelnone
	 */
	public static String getLabelNone() {
		return labelNone;
	}
}