#!/usr/bin/python3
# -*- coding: utf-8 -*-

from typing import List
import random

caesar_encoding: str = "ABCDEFGHIJKLMNOPQRSTUVWXYZ "


def key_gen() -> int:
    """
    Generates one Caesar key
    """
    # Correct
    key = random.randint(1, 26)

    # Bug
    # key = random.randint(0, 27)
    return key


def key_gen_test() -> bool:
    # What else could we test here?
    # What about randomness?
    # Will this test always work?
    # Does it guarantee correctness?
    for c in range(300):
        if (key_gen() > 26) or (key_gen() < 1):
            print("key_gen_test() failed")
            return False
    return True


def str_to_num_arr(message: str) -> List[int]:
    """
    Translates a string into a Caesar-encoded List
    """
    arr: List[int] = []
    for character in message:
        arr.append(caesar_encoding.find(character.upper()))
    return arr


def str_to_num_arr_test() -> bool:
    # TODO
    return True


def num_arr_to_str(encoded_arr: List[int]) -> str:
    """
    Translates a Caesar encoded list back into a string
    """
    plaintext: List[str] = []
    for counter, encoded_char in enumerate(encoded_arr):
        plaintext.append(caesar_encoding[encoded_char])
    plaintext_string = "".join(plaintext)
    return plaintext_string


def num_arr_to_str_test() -> bool:
    # TODO
    return True


def translate(encoded_arr: List[int], mode: int, key: int) -> List[int]:
    """
    Encrypts or decryps a Caesar-encoded List of ints
    """
    translated: List[int] = []
    for encoded_char in encoded_arr:
        if mode == 1:
            # 27 is the symbol set size (# letters in alphabet)
            # Note the space added above (bug from last time)!
            translated.append((encoded_char + key) % 27)
        else:
            translated.append((encoded_char - key) % 27)
    return translated


def translate_test() -> bool:
    # TODO
    return True


def encrypt_test() -> bool:
    for character in caesar_encoding:
        for key in range(1, 26):
            # Encrypt
            message_arr = str_to_num_arr(character)
            message_arr = translate(message_arr, 1, key)
            message = num_arr_to_str(message_arr)

            # Decrypt
            message_arr = str_to_num_arr(message)
            message_arr = translate(message_arr, 0, key)
            message = num_arr_to_str(message_arr)

            # check
            if message != character:
                print("encrypt_test() failed")
                return False
    return True


def run_tests() -> bool:
    """
    Note: greedy quitting with booleans means that
    you should put the more aggregate/complicated tests last,
    and the simple tests first.
    """
    return (
        key_gen_test()
        and str_to_num_arr_test()
        and num_arr_to_str_test()
        and translate_test()
        and encrypt_test()
    )


# Only runs main if tests pass:
if __name__ == "__main__" and run_tests():
    message: str = input("\nEnter your message, in English:\n")
    gen_key: str = input("Want to generate a key? (y/n)")

    if gen_key == "y":
        ok: int = 0
        while ok == 0:
            key = key_gen()
            print("Is the key it ok with you (1-yes, 0-no, make another): ")
            ok = int(input())
    else:
        key = int(input("What is your key (0-25)?"))

    print("\nYour Caesar key is: ")
    print(key)
    print("\n Share this with your partner. Don't tell anyone else\n")

    print("\nEnter 1 for encryption, and 0 for decryption: ")
    mode: int = int(input())
    message_arr = str_to_num_arr(message)
    message_arr = translate(message_arr, mode, key)
    message = num_arr_to_str(message_arr)
    print(message)