#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
How could this code be improved?
@author: taylor@mst.edu
"""

import secrets

LETTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ ."


def keygen(keylen: int) -> str:
    return "".join([secrets.choice(LETTERS) for _ in range(keylen)])


def main() -> None:
    myMessage = (
        "The generation of random numbers is too important to be left to chance."
    )
    myMode = "encrypt"  # Set to either 'encrypt' or 'decrypt'.
    if myMode == "encrypt":
        myKey = keygen(len(myMessage))
        print("Your ONE-TIME use key is:\n", myKey, "\n")
        translated = translateMessage(myKey, myMessage, "encrypt")
    elif myMode == "decrypt":
        myKey = input("What was the ONE-TIME use key you securily pre-shared?")
        translated = translateMessage(myKey, myMessage, "decrypt")
    print("Your " + myMode + "ed message was:", translated)


def translateMessage(key: str, message: str, mode: str) -> str:
    translated = []  # Stores the encrypted/decrypted message string.
    keyIndex = 0
    key = key.upper()
    for symbol in message:  # Loop through each symbol in message.
        num = LETTERS.find(symbol.upper())
        if num != -1:  # -1 means symbol.upper() was not found in LETTERS.
            if mode == "encrypt":
                num += LETTERS.find(key[keyIndex])  # Add if encrypting.
            elif mode == "decrypt":
                num -= LETTERS.find(key[keyIndex])  # Subtract if decrypting.
            num %= len(LETTERS)  # Handle any wraparound.
            # Add the encrypted/decrypted symbol to the end of translated:
            if symbol.isupper():
                translated.append(LETTERS[num])
            elif symbol.islower():
                translated.append(LETTERS[num].lower())
            keyIndex += 1  # Move to the next letter in the key.
        else:
            # Append the symbol without encrypting/decrypting.
            translated.append(symbol)
    return "".join(translated)


if __name__ == "__main__":
    main()