#! python
import os
import re

## Initialize global variables
continue_cipher = ""
demo_alphabet = "KRYPTOSABCDEFGHIJLMNQUVWXZ"
demo_key = "PALIMPSEST"
demo_cipher_string = "EMUFPHZLRFAXYUSDJKZLDKRNSHGNFIVJYQTQUXQBQVYUVLLTREVJYQTMKYRDMFD"
demo_cipher_decoded = "BETWEENSUBTLESHADINGANDTHEABSENCEOFLIGHTLIESTHENUANCEOFIQLUSION"


## Visuals
def display_header():
    print("################################################")
    print("#                                              #")
    print("#            --- VIGENERE CIPHER ---           #")
    print("#                                              #")
    print("#   A simple Vigenere cipher decoder/encoder   #")
    print("#                                              #")
    print("################################################", end="\n\n")
    return
def display_results(mode, cipher_vars):
    # Clear screen for final results
    os.system('cls')

    # Display header
    display_header()

    # Decompose cipher_vars
    (alphabet, key, cipher_string, results) = cipher_vars

    print("Mode:", "Decrypt" if mode == "D" else "Encrypt", end="\n\n")
    print("Alphabet:", alphabet)
    print("Key:", key)
    print("Cipher String:", cipher_string, end="\n\n")
    print("Decoded string:" if mode == "D" else "Encoded string:", results, end="\n\n")
    return


## Validations
def string_is_alpha(input_string):
    return True if re.match("^[a-zA-Z_]*$", input_string) else False


## Cipher variables
def get_alphabet():
    global demo_alphabet

    while True:
        alphabet = input("Enter cipher alphabet: ").upper()
        if alphabet == "":
            alphabet = demo_alphabet
            break
        elif string_is_alpha(alphabet) is False:
            print("The alphabet is not valid. Alphabet should not contain spaces, digits or special characters.")
        else:
            break

    return alphabet
def get_key():
    global demo_key

    while True:
        key = input("Enter cipher key: ").upper()
        if key == "":
            key = demo_key
            break
        elif string_is_alpha(key) is False:
            print("The key is not valid. Key should not contain spaces, digits or special characters.")
        else:
            break

    return key
def get_cipher_string(mode):
    global demo_cipher_string
    global demo_cipher_decoded

    while True:
        cipher_string = input("Enter cipher string: ").upper()
        if cipher_string == "":
            cipher_string = demo_cipher_string if mode == "D" else demo_cipher_decoded
            break
        elif string_is_alpha(cipher_string) is False:
            print("The cipher string is not valid. Cipher strings should not contain spaces, digits or special characters.")
        else:
            break

    return cipher_string


## Cipher actions
def get_cipher_alphabets(alphabet, key):
    cipher_alphabets = []

    for char in key:
        char_index = alphabet.find(char)
        cipher_alphabet = alphabet[char_index:] + alphabet[:char_index]
        cipher_alphabets.append(cipher_alphabet)

    return cipher_alphabets
def start_cipher(mode, alphabet, key, cipher_string):
    mode_string = ""
    cipher_alphabets = get_cipher_alphabets(alphabet, key)
    
    cipher_alphabet_index = 0
    for char in cipher_string:
        # Reset cipher_alphabet_index to 0 when at end of cipher alphabets
        if cipher_alphabet_index == len(cipher_alphabets):
            cipher_alphabet_index = 0
        
        # Use appropriate alphabet based on mode
        # Syntax: base_alphabet[mode_alphabet.find(char)]
        if mode == "D":
            mode_string += alphabet[cipher_alphabets[cipher_alphabet_index].find(char)]
        else:
            mode_string += cipher_alphabets[cipher_alphabet_index][alphabet.find(char)]

        cipher_alphabet_index += 1

    return mode_string


## Cipher Mode
def get_cipher_mode():
    while True:
        cipher_mode = input("Choose cipher mode - [D]ecrypt or [E]ncrypt: ").upper()
        if cipher_mode != "D" and cipher_mode != "E":
            print("That is not a valid option. Please enter 'D' for decrypt and 'E' for encrypt.")
        else:
            break

    print("")
    return cipher_mode
def start_cipher_mode(mode):
    print("Press 'enter' to use demo options")
    alphabet = get_alphabet()
    key = get_key()
    cipher_string = get_cipher_string(mode)
    mode_string = start_cipher(mode, alphabet, key, cipher_string)
    return alphabet, key, cipher_string, mode_string


## Loop cipher
def get_continue_cipher():
    while True:
        continue_cipher = input("Do you want to decode/encode more? [Y/N]: ").upper()
        if continue_cipher != "Y" and continue_cipher != "N":
            print("That is not a valid option. Please enter 'Y' to continue and 'N' to quit.")
        else:
            break
    return continue_cipher


## Start vigenere cipher program
while continue_cipher != "N":
    # Clear the screen after each operation
    os.system('cls')

    # Display header
    display_header()

    # Determine cipher mode
    cipher_mode = get_cipher_mode()
    cipher_vars = start_cipher_mode(cipher_mode)

    # Display results
    display_results(cipher_mode, cipher_vars)

    continue_cipher = get_continue_cipher()