#!/usr/bin/env python
import os,sys

ORD_a = ord('a') # 97
ORD_A = ord('A') # 65

def shift_cipher_encrypt(plaintext, key):
    # apply encryption to text with given key
    encrypted_message = ''

    for char in plaintext:
        if char.isalpha():
            ascii_offset = ORD_A if char.isupper() else ORD_a  # Determine ASCII offset based on uppercase or lowercase letter

            # the comment shown below are the pseudo code, it demonstrate the ideas only
            # let say the input is 'the' // without quote
            # find distance of target character with reference to A or a
            # i.e. t - a = 19 ,   h - a = 7 ,  e - a = 4
            distance = ord(char) - ascii_offset

            # [19,7,4] + [8,8,8] (key) = [27, 15, 12]
            # Shift the character by adding the key and taking modulo 26 to wrap around
            # [27,15,12] % [26,26,26] = [1,15,12]   // get modules
            shifted_distance = (distance + key) % 26

            # [1,15,12] + [97,97,97] = [98,112,109]
            # chr(98) , chr(112) , chr(109) = 'bpm'
            shifted_char = chr(shifted_distance + ascii_offset)

            # so: the -> bpm
            encrypted_message += shifted_char

        else:
            # consider integer case, retain
            encrypted_message += char

    return encrypted_message


def shift_cipher_decrypt(ciphertext, key):
    plaintext = ""

    for char in ciphertext:
        if char.isalpha():
            ascii_offset = ORD_a if char.islower() else ORD_A  # Determine ASCII offset based on lowercase or uppercase letter

            # Calculate the distance of the target character from a or A
            distance = ord(char) - ascii_offset

            # apply shift, get the remainder of 26
            shifted_distance = (distance - key) % 26

            # Convert back to ASCII
            decrypted_char = chr(shifted_distance + ascii_offset)

            plaintext += decrypted_char
        else:
            # If it is not an alphabetic character, retain as is.
            plaintext += char

    return plaintext


def count_letter_e(txt_in):
    # reserved function for demonstration purpose
    occurence = 0
    for char in txt_in:
        if char.isalpha():
            if char.lower() == 'e':
                occurence += 1
    return occurence

def count_most_occurrence_letter(txt_in):
    # letter e, as stated have the most occurrence in the message by statistics.
    # as 'Shift Cipher' is a encryption by letter shifting, the letters have good chance
    # to have the most occurrence too in the encrypted text.
    output = [0] * 26    # bucket for 26 letters

    for char in txt_in:
        if char.isalpha():
            output[ord(char.lower()) - ORD_a] += 1

    # output contains the statistics of paragraph letter by letter
    return output

def find_max_occurrence(char_occurrences):
    # get the letter of the most occurrences. i.e. m
    # by subtract between this letter to e, k can be guess

    # find max occurrence and its index
    max_idx = char_occurrences.index(max(char_occurrences))

    # subtract it with index of e -> 4
    return max_idx - 4

def encrypt_file(file_path, key=8):
    # open a file and apply encryption
    output_file = file_path.replace('.txt','_e.txt')

    # convert it to integer
    key = int(key)

    # open source file (plaintext)
    with open(file_path,'r') as fi:
        temp = ''.join(fi.readlines())

        # open target file (encrypted text)
        with open(output_file,'r+') as fo:
            fo.truncate(0)
            fo.writelines([shift_cipher_encrypt(temp, key)])

    print(f'encryption done and file saved to {output_file}')
    return

def decrypt_file(file_path):
    # will open an encrypted file and decrypt it by a guessed key

    with open(file_path,'r') as fi:
        # beginning of the process
        # read file and join the lines all
        lines = fi.readlines()
        e_temp = ''.join(lines)

        characters_distribution = count_most_occurrence_letter(e_temp)

        print('')
        print('distribution of letters in encrypted text (case insensitive, from a to z)')
        print(characters_distribution)

        print('')
        guess_k = find_max_occurrence(characters_distribution)
        print(f'guessed k: {guess_k}')

        print('')
        print('decrypted text:')
        decrypted_text = shift_cipher_decrypt(e_temp, guess_k)
        print(decrypted_text)

while True:
    # show menu
    print()
    print("1. Encrypt File")
    print("2. Decrypt File")
    print("q. quit")
    print()
    option = input("Select an option (1/2/q): ")

    if option == "1":
        # run if user want to encrypt file
        # check if user entered a file
        user_not_enter_file = True
        while user_not_enter_file:
          file_path = input("Enter the path of the file to encrypt: ")
          if len(file_path) > 0:
            if os.path.exists(file_path):
              user_not_enter_file = False
            else:
              print('sorry but the file not exist')
          else:
              print('please enter a file path')

        # check if user entered a key
        user_not_enter_key = True
        while user_not_enter_key:
          key = input("Enter the key(k) to encrypt: ")
          if (len(key) > 0):
              user_not_enter_key = False
          else:
              print('please enter a key(k)')

        if os.path.exists(file_path):
            encrypt_file(file_path, key)
            print('encryption done')
        else:
            print("File does not exist.")

    elif option == "2":
        # run if user want to decrypt file
        file_path = input("Enter the path of the file to decrypt: ")
        if os.path.exists(file_path):
            decrypt_file(file_path)
            print('decryption done')
        else:
            print("File does not exist.")

    elif option.lower() == "q":
        print('quitting bye ...')
        break
    else:
        print('')
        print('ERROR !')
        print('please enter either [1/2/q]')
        input("press a key to continue ...")
        print('')

print("Exiting...")