#!/usr/bin/env python3
import argparse
import string

parser: argparse.ArgumentParser = argparse.ArgumentParser(
    prog='Basic Vigenère Cypher',
    description='A quick and dirty implementation of the Vigenére Cypher. Accepts any key as the key.',
    epilog='Enter the message first and the key second.'
)

group = parser.add_mutually_exclusive_group(required=True)
group.add_argument('-e', '--encrypt', nargs=2, help="Message to be decrypted or encrypted and the key used")
group.add_argument('-d', '--decrypt', nargs=2, help="Message to be decrypted or encrypted and the key used")

class Vigenere:
    def __init__(self) -> None:
        self.args = parser.parse_args()
        self.chars: str = ''.join(c for c in string.printable if c.isprintable()) + ''.join(chr(c) for c in range(128, 256) if not chr(c).isspace or c in (160,)) # All Characters from extended ASCII, excluding control codes
        if self.args.encrypt:
            self.msg: str = self.args.encrypt[0]
            self.key: str = self._generate_key(self.args.encrypt[1])
        else:
            self.msg: str = self.args.decrypt[0]
            self.key: str = self._generate_key(self.args.decrypt[1])

    def _generate_key(self, key: str) -> str:
        if len(key) < len(self.msg):
            key = (key * (len(self.msg) // len(key) + 1))[:len(self.msg)]
        else:
            key = key[:len(self.msg)]
        return key

    def _create_alphabet(self, shift: int) -> str:
        alphabet: str = self.chars
        shift = (shift % len(alphabet))
        return alphabet[shift:] + alphabet[:shift]

    def _find_index(self, index: int) -> int:
        return self.chars.index(self.key[index])

    def encrypt(self) -> str:
        encrypted_msg: str = ""
        for index in range(len(self.msg)):
            char: str = self.msg[index]
            lookup_alphabet: str = self._create_alphabet(self._find_index(index))
            original_index: int = self.chars.index(char)
            encrypted_msg += lookup_alphabet[original_index]

        return encrypted_msg

    def decrypt(self) -> str:
        decrypted_msg: str = ""
        for index in range(len(self.msg)):
            lookup_alphabet: str = self._create_alphabet(self._find_index(index))
            location_index: int = lookup_alphabet.index(self.msg[index])
            decrypted_msg += self.chars[location_index]

        return decrypted_msg

    def output(self) -> None:
        if self.args.encrypt:
            ciphertext = self.encrypt()
            print(ciphertext)
        else:
            plaintext = self.decrypt()
            print(plaintext)

if __name__ == "__main__":
    vigenere_cipher = Vigenere()
    vigenere_cipher.output()