import static java.lang.System.err;
import static java.lang.System.out;

import java.io.ByteArrayOutputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Arrays;
import java.util.Base64;
import javax.crypto.Cipher;
import javax.crypto.Mac;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;

class aes {
	private static final byte[] MAGIC = new byte[] { 'D', 'F', 'A', 'E', 'S', '1' };
	private static final int KEY_BYTES = 64;
	private static final int AES_KEY_BYTES = 32;
	private static final int IV_BYTES = 12;
	private static final int GCM_TAG_BITS = 128;
	private static final int HEADER_BYTES = MAGIC.length + IV_BYTES;

	enum Command {
		ENCRYPT, DECRYPT, MERGE
	}

	public static void main(String[] arguments) throws Exception {
		try {
			var exit = run(arguments);
			if (exit != 0) {
				System.exit(exit);
			}
		} catch (IllegalArgumentException problem) {
			err.println(problem.getMessage());
			System.exit(2);
		} catch (Exception problem) {
			err.println(problem.getMessage());
			System.exit(1);
		}
	}

	private static int run(String[] arguments) throws Exception {
		if (arguments.length == 0) {
			usage();
			throw new IllegalArgumentException("missing command");
		}

		var command = parseCommand(arguments[0]);
		var wrongArgumentCount = command == Command.MERGE ? arguments.length != 6 : arguments.length != 3 && arguments.length != 5;
		if (wrongArgumentCount) {
			usage();
			throw new IllegalArgumentException("wrong argument count");
		}
		if (!"--key-file".equals(arguments[1])) {
			usage();
			throw new IllegalArgumentException("expected --key-file");
		}

		var key = readKey(Path.of(arguments[2]));
		if (command == Command.MERGE) {
			return merge(key, Path.of(arguments[3]), Path.of(arguments[4]), Path.of(arguments[5]));
		}

		var input = arguments.length == 3 ? System.in.readAllBytes() : Files.readAllBytes(Path.of(arguments[3]));
		var output = command == Command.ENCRYPT ? encrypt(key, input) : decrypt(key, input);
		if (arguments.length == 3) {
			System.out.write(output);
		} else {
			var outputPath = Path.of(arguments[4]);
			if (Files.exists(outputPath)) {
				throw new IllegalArgumentException("output path already exists: " + outputPath);
			}
			try (var stream = Files.newOutputStream(outputPath)) {
				stream.write(output);
			}
		}
		return 0;
	}

	static int merge(byte[] key, Path base, Path current, Path other) throws Exception {
		var temporary = Files.createTempDirectory("dotfiles-aes-merge-");
		var plainBase = temporary.resolve("base");
		var plainCurrent = temporary.resolve("current");
		var plainOther = temporary.resolve("other");
		try {
			Files.write(plainBase, mergeInput(key, base));
			Files.write(plainCurrent, mergeInput(key, current));
			Files.write(plainOther, mergeInput(key, other));
			var process = new ProcessBuilder("git", "merge-file", "-L", "current", "-L", "base", "-L", "other",
					plainCurrent.toString(), plainBase.toString(), plainOther.toString()).inheritIO().start();
			var exit = process.waitFor();
			if (exit > 1) {
				throw new IllegalStateException("git merge-file failed with exit " + exit);
			}
			Files.write(current, encrypt(key, Files.readAllBytes(plainCurrent)));
			return exit;
		} finally {
			Files.deleteIfExists(plainBase);
			Files.deleteIfExists(plainCurrent);
			Files.deleteIfExists(plainOther);
			Files.deleteIfExists(temporary);
		}
	}

	private static byte[] mergeInput(byte[] key, Path path) throws Exception {
		var input = Files.readAllBytes(path);
		return input.length == 0 ? input : decrypt(key, input);
	}

	static byte[] encrypt(byte[] key, byte[] plaintext) throws Exception {
		if (startsWith(plaintext, MAGIC)) {
			throw new IllegalArgumentException("input is already a dotfiles AES blob");
		}

		var split = splitKey(key);
		var iv = deterministicIv(split.macKey(), plaintext);
		var header = header(iv);
		var cipher = Cipher.getInstance("AES/GCM/NoPadding");
		cipher.init(Cipher.ENCRYPT_MODE, new SecretKeySpec(split.encryptionKey(), "AES"), new GCMParameterSpec(GCM_TAG_BITS, iv));
		cipher.updateAAD(header);
		var ciphertext = cipher.doFinal(plaintext);

		var output = new ByteArrayOutputStream(header.length + ciphertext.length);
		output.write(header);
		output.write(ciphertext);
		return output.toByteArray();
	}

	static byte[] decrypt(byte[] key, byte[] encrypted) throws Exception {
		if (!startsWith(encrypted, MAGIC)) {
			throw new IllegalArgumentException("input is not a dotfiles AES blob");
		}
		if (encrypted.length < HEADER_BYTES + 1) {
			throw new IllegalArgumentException("encrypted blob is truncated");
		}

		var iv = Arrays.copyOfRange(encrypted, MAGIC.length, HEADER_BYTES);
		var header = Arrays.copyOfRange(encrypted, 0, HEADER_BYTES);
		var ciphertext = Arrays.copyOfRange(encrypted, HEADER_BYTES, encrypted.length);
		var split = splitKey(key);
		var cipher = Cipher.getInstance("AES/GCM/NoPadding");
		cipher.init(Cipher.DECRYPT_MODE, new SecretKeySpec(split.encryptionKey(), "AES"), new GCMParameterSpec(GCM_TAG_BITS, iv));
		cipher.updateAAD(header);
		return cipher.doFinal(ciphertext);
	}

	static boolean startsWith(byte[] bytes, byte[] prefix) {
		if (bytes.length < prefix.length) {
			return false;
		}
		for (var index = 0; index < prefix.length; index++) {
			if (bytes[index] != prefix[index]) {
				return false;
			}
		}
		return true;
	}

	private static Command parseCommand(String input) {
		if ("encrypt".startsWith(input)) {
			return Command.ENCRYPT;
		}
		if ("decrypt".startsWith(input)) {
			return Command.DECRYPT;
		}
		if ("merge".startsWith(input)) {
			return Command.MERGE;
		}
		throw new IllegalArgumentException("unknown command: " + input);
	}

	private static byte[] readKey(Path path) throws Exception {
		if (!Files.exists(path)) {
			throw new IllegalArgumentException("key file does not exist: " + path);
		}
		var text = Files.readString(path).trim();
		var key = Base64.getDecoder().decode(text);
		if (key.length != KEY_BYTES) {
			throw new IllegalArgumentException("key file must contain " + KEY_BYTES + " base64-encoded bytes");
		}
		return key;
	}

	private static SplitKey splitKey(byte[] key) {
		return new SplitKey(
				Arrays.copyOfRange(key, 0, AES_KEY_BYTES),
				Arrays.copyOfRange(key, AES_KEY_BYTES, KEY_BYTES));
	}

	private static byte[] deterministicIv(byte[] macKey, byte[] plaintext) throws Exception {
		var mac = Mac.getInstance("HmacSHA256");
		mac.init(new SecretKeySpec(macKey, "HmacSHA256"));
		return Arrays.copyOf(mac.doFinal(plaintext), IV_BYTES);
	}

	private static byte[] header(byte[] iv) throws Exception {
		var output = new ByteArrayOutputStream(HEADER_BYTES);
		output.write(MAGIC);
		output.write(iv);
		return output.toByteArray();
	}

	private static void usage() {
		out.println("Symmetrically encrypts/decrypts dotfiles data.");
		out.println("Usage for streams: java tools/aes.java <encrypt|decrypt> --key-file <path>");
		out.println("Usage for files:   java tools/aes.java <encrypt|decrypt> --key-file <path> <input> <output>");
		out.println("Usage for merges:  java tools/aes.java merge --key-file <path> <base> <current> <other>");
	}

	static class SplitKey {
		private final byte[] encryptionKey;
		private final byte[] macKey;

		SplitKey(byte[] encryptionKey, byte[] macKey) {
			this.encryptionKey = encryptionKey;
			this.macKey = macKey;
		}

		byte[] encryptionKey() {
			return encryptionKey;
		}

		byte[] macKey() {
			return macKey;
		}
	}
}