import sys
from math import ceil

FREQ_ENG = {'e': 12.7, 't': 9.06, 'a': 8.17, 'o': 7.51, 'i': 6.97, 'n': 6.75, 's': 6.33,
            'h': 6.09, 'r': 5.99, 'd': 4.25, 'l': 4.03, 'c': 2.78, 'u': 2.76, 'm': 2.41,
            'w': 2.36, 'f': 2.23, 'g': 2.02, 'y': 1.97, 'p': 1.93, 'b': 1.49, 'v': 0.98,
            'k': 0.77, 'x': 0.15, 'j': 0.15, 'q': 0.1, 'z': 0.05}


# ALPHABET = ''' '"#$%&()*+,-./0123456789:;<=>?@[]^_`abcdefghijklmnopqrstuvwxyz{|}~!'''
# ALPHABET = '''abcdefghijklmnopqrstuvwxyz '.,()0123456789'''
ALPHABET = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
M = len(ALPHABET)


def main(args):
    file_input, file_output, mode, key = args[:4]
    output = str
    with open(file_input) as f:
        text = f.read()
    if mode == 'encrypt':
        output = encrypt(text, key)
    elif mode == 'decrypt':
        output = decrypt(text, key)
    print('Input text:\n{}\n'.format(text))
    print('Output text:\n{}\n'.format(output))
    with open(file_output, 'w') as f:
        f.write(output)


def encrypt(text, key):
    key = (key * ceil(len(text) / len(key)))[:len(text)]
    result = []
    for i, symbol in enumerate(text):
        if symbol.upper() in ALPHABET:
            if symbol.isupper():
                result.append(ALPHABET[(ALPHABET.find(symbol)+ALPHABET.find(key[i])) % M])
            else:
                result.append(ALPHABET[(ALPHABET.find(symbol.upper()) + ALPHABET.find(key[i])) % M].lower())
        else:
            result.append(symbol)
    # result = [ALPHABET[(ALPHABET.find(text[i].upper())+ALPHABET.find(key[i])) % M] if text[i] in ALPHABET else
    #           text[i] for i in range(len(text))]
    return ''.join(result)


def encrypt_one(message, key, mode='encrypt'):
    translated = []
    for symbol in message: # loop through each character in message
        num = ALPHABET.find(symbol.upper())
        if num != -1: # -1 means symbol.upper() was not found in ALPHABET
            if mode == 'encrypt':
                num += ALPHABET.find(key[keyIndex]) # add if encrypting
            elif mode == 'decrypt':
                num -= ALPHABET.find(key[keyIndex]) # subtract if decrypting

            num %= len(ALPHABET) # handle the potential wrap-around

            # add the encrypted/decrypted symbol to the end of translated.
            if symbol.isupper():
                translated.append(ALPHABET[num])
            elif symbol.islower():
                translated.append(ALPHABET[num].lower())

            keyIndex += 1 # move to the next letter in the key
            if keyIndex == len(key):
                keyIndex = 0
        else:
            # The symbol was not in ALPHABET, so add it to translated as is.
            translated.append(symbol)


def decrypt(text, key):
    key = (key * ceil(len(text) / len(key)))[:len(text)]
    result = [ALPHABET[(ALPHABET.find(text[i])-ALPHABET.find(key[i])) % M] if text[i] in ALPHABET else
              text[i] for i in range(len(text))]
    return ''.join(result)


if __name__ == '__main__':
    # try:

    # main(sys.argv[1:])
    main(('1.txt', '2.txt', 'encrypt', 'TEST'))
    # except ValueError:
    #     exit('''
    # Формат ввода данных:
    # input_file output_file mode key
    # ''')
    # except KeyboardInterrupt:
    #     exit('''
    # Работа прервана
    #         ''')