"""
Szyfr książkowy wykorzystujący powieść Zaginiony świat

Słowa nieznajdujące się w książce zostają przeliterowane 
z użyciem pierwszych liter słów. Początek trybu literowania 
to 'a a', a koniec - 'the the'.

autor: Eric T. Mortenson
"""
import os
import random
import string
import sys
from collections import defaultdict, Counter


def main():
    message = input("Wpisz szyfrogram lub tekst jawny: ") 
    process = input("Wpisz 'encrypt' (szyfrowanie) "
                    "lub 'decrypt' (odszyfrowywanie): ")
    shift = int(input("Przesuniecie indeksu (1-365) = "))
    infile = input("Wpisz nazwe pliku wraz z rozszerzeniem: ")

    if not os.path.exists(infile):
        print("Nie znaleziono pliku {}. Program konczy dzialanie."
              .format(infile), file=sys.stderr)
        sys.exit(1)
    word_list = load_file(infile)
    word_dict = make_dict(word_list, shift)
    letter_dict = make_letter_dict(word_list)

    if process == 'encrypt':
        ciphertext = encrypt(message, word_dict, letter_dict)
        encryptedWordList = []
        for number in ciphertext:
            encryptedWordList.append(word_list[number - shift])

        print("\nlista zaszyfrowanych wyrazow = \n {}\n"
              .format(' '.join(encryptedWordList)))
        print("szyfrogram = \n {}\n".format(ciphertext))

        # Sprawdza szyfrowanie poprzez odszyfrowanie szyfrogramu.
        print("tekst jawny = ")
        singleFirstCheck = False
        for cnt, i in enumerate(ciphertext):
            if word_list[ciphertext[cnt] - shift] == 'a' and \
                    word_list[ciphertext[cnt + 1] - shift] == 'a':
                continue
            if word_list[ciphertext[cnt] - shift] == 'a' and \
                    word_list[ciphertext[cnt - 1] - shift] == 'a':
                singleFirstCheck = True
                continue
            if singleFirstCheck == True and cnt < len(ciphertext) - 1 and \
                    word_list[ciphertext[cnt] - shift] == 'the' and \
                    word_list[ciphertext[cnt + 1] - shift] == 'the':
                continue
            if singleFirstCheck == True and \
                    word_list[ciphertext[cnt] - shift] == 'the' and \
                    word_list[ciphertext[cnt - 1] - shift] == 'the':
                singleFirstCheck = False
                print(' ', end='', flush=True)
                continue
            if singleFirstCheck == True:
                print(word_list[i - shift][0], end='', flush=True)
            if singleFirstCheck == False:
                print(word_list[i - shift], end=' ', flush=True)

    elif process == 'decrypt':
        plaintext = decrypt(message, word_list, shift)
        print("\nszyfrogram = \n{}".format(plaintext))


def load_file(infile):
    """
    Odczytuje i zwraca plik tekstowy jako listę słów 
    złożonych z małych liter.
    """
    with open(infile, encoding='utf-8') as file:
        words = [word.lower() for line in file for word in line.split()]
        words_no_punct = ["".join(char for char in word if char not in \
                                  string.punctuation) for word in words]
    return words_no_punct


def make_dict(word_list, shift):
    """
    Zwraca słownik znaków jako kluczy i przesuniętych indeksów 
    jako wartości.
    """
    word_dict = defaultdict(list)
    for index, word in enumerate(word_list):
        word_dict[word].append(index + shift)
    return word_dict


def make_letter_dict(word_list):
    firstLetterDict = defaultdict(list)
    for word in word_list:
        if len(word) > 0:
            if word[0].isalpha():
                firstLetterDict[word[0]].append(word)
    return firstLetterDict


def encrypt(message, word_dict, letter_dict):
    """Zwraca listę indeksów reprezentujących znaki w wiadomości."""
    encrypted = []
    # Usuwa znaki interpunkcyjne ze słów wiadomości.
    messageWords = message.lower().split()
    messageWordsNoPunct = ["".join(char for char in word if char not in \
                                   string.punctuation) for word in
                           messageWords]
    for word in messageWordsNoPunct:
        if len(word_dict[word]) > 1:
            index = random.choice(word_dict[word])
        # Random.choice kończy się błędem, jeżeli jest tylko 1 opcja do wyboru.
        elif len(word_dict[word]) == 1:
            index = word_dict[word][0]
        # Słowo nie znajduje się w word_dict.
        elif len(word_dict[word]) == 0:
            encrypted.append(random.choice(word_dict['a']))
            encrypted.append(random.choice(word_dict['a']))

            for letter in word:
                if letter not in letter_dict.keys():
                    print('\nBrak znaku {} w slowniku.'
                        .format(letter), file=sys.stderr)
                    continue
                if len(letter_dict[letter]) > 1:
                    newWord = random.choice(letter_dict[letter])
                else:
                    newWord = letter_dict[letter][0]
                if len(word_dict[newWord]) > 1:
                    index = random.choice(word_dict[newWord])
                else:
                    index = word_dict[newWord][0]
                encrypted.append(index)

            encrypted.append(random.choice(word_dict['the']))
            encrypted.append(random.choice(word_dict['the']))
            continue
        encrypted.append(index)
    return encrypted


def decrypt(message, word_list, shift):
    """
    Odszyfrowuje szyfrogram i zwraca tekst jawny.
    Pokazuje, jak wygląda tekst przed wyciągnięciem pierwszych liter.
    """
    plaintextList = []
    indexes = [s.replace(',', '').replace('[', '').replace(']', '')
               for s in message.split()]
    for count, i in enumerate(indexes):
        plaintextList.append(word_list[int(i) - shift])
    return ' '.join(plaintextList)


def check_for_fail(ciphertext):
    """Zwraca True, jeżeli szyfrogram zawiera powielone klucze."""
    check = [k for k, v in Counter(ciphertext).items() if v > 1]
    if len(check) > 0:
        print(check)
        return True


if __name__ == '__main__':
    main()
