import binascii
import string
from typing import List, Dict, Set


def load_ciphers() -> List[bytes]:
    with open("data/19_enc.txt", "r") as f:
        r = [binascii.a2b_base64(line) for line in f]
    return r


def xor_lookup_set(letters: str) -> Dict[int, Set[str]]:
    d = {i: set() for i in range(256)}
    for c1 in letters:
        for c2 in letters:
            xored = ord(c1) ^ ord(c2)
            d[xored].add(c1)
            d[xored].add(c2)
    return d


def attack(ciphers: List[bytes]) -> List[List[str]]:
    """ Find out possible characters for each cipher pair and hope that there
    is only one possible character. If no character was found add '?' and
    if more than one was found add '^'. """

    ciphers_len = len(ciphers)
    deciphered = [[] for _ in range(ciphers_len)]
    max_cipher_len = max(map(len, ciphers))
    for byte_index in range(0, max_cipher_len):
        # Certain bytes only work with certain letters (found empirically).
        if byte_index == 10:
            LETTERS = string.ascii_letters + " _-.,;:'"
        elif byte_index == 20:
            LETTERS = string.ascii_letters + " _-.,;:?"
        else:
            LETTERS = string.ascii_letters + " _-.,;:"
        lookup = xor_lookup_set(LETTERS)
        target_bytes = [cipher[byte_index]
                        if len(cipher) > byte_index
                        else None
                        for cipher in ciphers]
        possible_chars = [set(LETTERS) for _ in range(ciphers_len)]
        for i in range(ciphers_len):
            for j in range(i, ciphers_len):
                if target_bytes[i] is None or target_bytes[j] is None:
                    continue
                xored = target_bytes[i] ^ target_bytes[j]
                chars = lookup[xored]
                possible_chars[i] &= chars
                possible_chars[j] &= chars
        for cipher_index in range(ciphers_len):
            if len(ciphers[cipher_index]) <= byte_index:
                continue
            chars = list(possible_chars[cipher_index])
            match len(chars):
                case 0:
                    # print(f"No chars for {cipher_index=} {byte_index=}")
                    deciphered[cipher_index].append('?')
                case 1:
                    deciphered[cipher_index].append(chars[0].lower())
                case 2:
                    if chars[0].lower() == chars[1].lower():
                        deciphered[cipher_index].append(chars[0].lower())
                    else:
                        # print(f"Two {chars=} {cipher_index=} {byte_index=}")
                        deciphered[cipher_index].append('^')
                case _:
                    # print(f"Too many {chars=} {cipher_index=} {byte_index=}")
                    deciphered[cipher_index].append('^')
    return deciphered


def manual(decrypts: List[List[str]]) -> List[bytes]:
    """ Manually add guessed letters. """
    decrypts[0][30] = 'y'
    decrypts[2][30] = 'y'
    decrypts[4][30] = 'e'
    decrypts[4][32] = 'h'
    decrypts[4][33] = 'e'
    decrypts[4][34] = 'a'
    decrypts[4][35] = 'd'
    decrypts[6][30] = 'i'
    decrypts[13][30] = ' '
    decrypts[20][30] = ' '
    decrypts[25][30] = 'n'
    decrypts[28][30] = ' '
    decrypts[29][30] = 't'
    decrypts[37][30] = 'i'
    decrypts = list(map(lambda l: "".join(l), decrypts))
    return decrypts


if __name__ == "__main__":
    ciphers = load_ciphers()
    decrypts = manual(attack(ciphers))
    for d in decrypts:
        print(d)