import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map.Entry;

public final class Stegano {
	private final HashMap<Long, char[]> dict = new HashMap<>();

	public void train(String... texts) {
		dict.clear();
		var d = new HashMap<Long, HashMap<Character, Integer>>();
		for (var text : texts) {
			var k = -1L;
			for (int i = 0, n = text.length(); i < n; i++) {
				var c = text.charAt(i);
				incMapCounter(d, k, c);
				incMapCounter(d, k & 0xffffL, c);
				incMapCounter(d, k & 0xffff_ffffL, c);
				incMapCounter(d, k & 0xffff_ffff_ffffL, c);
				k = (k << 16) + c;
			}
		}
		int max = 0;
		for (var e : d.entrySet()) {
			var set = e.getValue().entrySet();
			int n = set.size();
			max = Math.max(max, n);
			@SuppressWarnings("unchecked")
			var es = (Entry<Character, Integer>[])set.toArray(new Entry[n]);
			Arrays.sort(es, (a, b) -> {
				var c = b.getValue() - a.getValue();
				return c != 0 ? c : a.getKey() - b.getKey();
			});
			var cs = new char[n];
			for (int i = 0; i < n; i++)
				cs[i] = es[i].getKey();
			dict.put(e.getKey(), cs);
		}
		if (max <= 1)
			throw new IllegalArgumentException("bad train texts");
	}

	private static void incMapCounter(HashMap<Long, HashMap<Character, Integer>> map, long k, char c) {
		map.computeIfAbsent(k, _ -> new HashMap<>()).compute(c, (_, n) -> n != null ? n + 1 : 1);
	}

	public String encrypt(byte[] src) {
		return encrypt(src, 0, src.length);
	}

	public String encrypt(byte[] src, int pos, int len) {
		if ((pos | len) < 0 || pos > src.length - len)
			throw new IllegalArgumentException();
		var sb = new StringBuilder((len + 4) << 2);
		var k = encrypt(sb, -1, new byte[]{(byte)len, (byte)(len >> 8), (byte)(len >> 16), (byte)(len >> 24)}, 0, 4);
		encrypt(sb, k, src, pos, pos + len);
		return sb.toString();
	}

	private long encrypt(StringBuilder sb, long k, byte[] src, int pos, int end) {
		var b = 0L;
		for (int s = 64; pos < end || s < 64; ) {
			var cs = findDict(k);
			char c;
			int n = cs.length;
			if (n > 1) {
				while (s >= 8 && pos < end)
					b += (src[pos++] & 0xffL) << (s -= 8);
				var r = decode(b, n);
				var bits = (int)(r >> 32);
				b <<= bits;
				s += bits;
				c = cs[(int)r];
			} else
				c = cs[0];
			k = (k << 16) + c;
			sb.append(c);
		}
		return k;
	}

	private char[] findDict(long k) {
		var cs = dict.get(k);
		if (cs == null || cs.length <= 1) {
			var cs1 = cs;
			cs = dict.get(k & 0xffff_ffff_ffffL);
			if (cs == null || cs.length <= 1) {
				if (cs1 == null)
					cs1 = cs;
				cs = dict.get(k & 0xffff_ffffL);
				if (cs == null || cs.length <= 1) {
					if (cs1 == null)
						cs1 = cs;
					cs = dict.get(k & 0xffffL);
					if (cs == null || cs.length <= 1) {
						if (cs1 == null)
							cs1 = cs;
						cs = dict.get(-1L);
						if (cs != null) {
							if (cs.length == 1 && cs1 != null)
								cs = cs1;
						} else if (cs1 != null)
							cs = cs1;
						else
							throw new IllegalStateException();
					}
				}
			}
		}
		return cs;
	}

	private static long decode(long b, int n) {
		assert n >= 2;
		int bits = 32 - Integer.numberOfLeadingZeros(n);
		int j = (1 << bits) - n;
		int i = (int)(b >>> (65 - bits));
		return i < j ? ((long)(bits - 1) << 32) + i : ((long)bits << 32) + ((b >>> (64 - bits)) - j);
	}

	private static long encode(int i, int n) {
		assert n >= 2 && (i & 0xffff_ffffL) < n;
		int bits = 32 - Integer.numberOfLeadingZeros(n);
		int j = (1 << bits) - n;
		return i < j ? ((long)(bits - 1) << 32) + i : ((long)bits << 32) + ((long)j << 1) + (i - j);
	}

	private long decrypt(String str, int[] pos, long k, byte[] dst) {
		var b = 0L;
		int s = 64;
		int p = pos[0];
		for (int j = 0, e = dst.length; j < e; ) {
			var c = str.charAt(p++);
			var cs = findDict(k);
			int i = 0, n = cs.length;
			while (cs[i] != c) {
				if (++i >= n)
					throw new IllegalStateException();
			}
			if (n > 1) {
				var r = encode(i, n);
				s -= (int)(r >> 32);
				b += (r & 0xffff_ffffL) << s;
				while (s <= 56 && j < e) {
					dst[j++] = (byte)(b >> 56);
					b <<= 8;
					s += 8;
				}
			}
			k = (k << 16) + c;
		}
		pos[0] = p;
		return k;
	}

	public byte[] decrypt(String str) {
		var b = new byte[4];
		var pos = new int[1];
		var k = decrypt(str, pos, -1, b);
		b = new byte[(b[0] & 0xff) + ((b[1] & 0xff) << 8) + ((b[2] & 0xff) << 16) + (b[3] << 24)];
		decrypt(str, pos, k, b);
		return b;
	}

	public static void main(String[] args) throws Exception {
		if (args.length == 0)
			System.out.println("usage: java Stegano <e|d> <train_utf8.txt> <input_file> <output_file>");
		else {
			var s = new Stegano();
			switch (args[0].charAt(0)) {
				case 'e':
					s.train(Files.readString(Paths.get(args[1]), StandardCharsets.UTF_8));
					var e = s.encrypt(Files.readAllBytes(Paths.get(args[2])));
					Files.writeString(Paths.get(args[3]), e, StandardCharsets.UTF_8);
					break;
				case 'd':
					s.train(Files.readString(Paths.get(args[1]), StandardCharsets.UTF_8));
					var d = s.decrypt(Files.readString(Paths.get(args[2]), StandardCharsets.UTF_8));
					Files.write(Paths.get(args[3]), d);
					break;
				default:
					s.train("ABCABA");
					var t = s.encrypt(new byte[]{1, 2});
					System.out.println(t);
					System.out.println(Arrays.toString(s.decrypt(t)));
					break;
			}
			System.out.println("done!");
		}
	}
}