import argparse

KEY = "2001"


def vigenere_cipher(text, key, decrypt=False):
    result = []
    key_length = len(key)
    key_as_int = [ord(i) for i in key]

    for i, char in enumerate(text):
        if char.isalpha():
            # Determine the shift
            key_index = i % key_length
            shift = key_as_int[key_index] - ord('A')

            # Convert the character to number (0-25), apply the shift, and convert back
            if char.isupper():
                base = ord('A')
            else:
                base = ord('a')

            if not decrypt:
                shifted = (ord(char) - base + shift) % 26
            else:  # decrypt
                shifted = (ord(char) - base - shift) % 26

            result.append(chr(shifted + base))
        else:
            result.append(char)

    return ''.join(result)


def main():
    parser = argparse.ArgumentParser(description="Vigenère Cipher Encryption/Decryption")
    parser.add_argument("message", help="Message to encrypt or decrypt")
    parser.add_argument("--decrypt", action="store_true", help="Decrypt the message")

    args = parser.parse_args()

    result = vigenere_cipher(args.message, KEY, args.decrypt)

    print(result)


if __name__ == "__main__":
    main()
