import os
import sys
from cryptography.fernet import Fernet

FILE_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)))

class Encryption:
    def __init__(self, op):
        self.set_key(op)
        self.run(op)

    def set_key(self, op):
        """
        Function to set encryption/decryption key.
        If no encryption key is given, it is generated by default.
        
        :param op: encrypt or decrypt
        :raises Exception: raises an exception if no decryption key is given.
        """
        self.keys = dict()
        if 'key.txt' in os.listdir(os.path.join(FILE_PATH, 'input')):
            key = self.read_file(os.path.join(FILE_PATH, 'input', 'key.txt'))
            self.key = Fernet(key=key)
            print(f'Sucessfully loaded {op}ion key')
        elif os.listdir(os.path.join(FILE_PATH, 'input')) and op == 'decrypt':
            raise Exception('Please provide key.txt for the files to be decrypted')
        elif op == 'encrypt':
            key = Fernet.generate_key()
            self.key = Fernet(key=key)
            self.write_output('key.txt', key)
            print('Sucessfully generated key')

    def read_file(self, file):
        """
        Function to read the file as bytes

        :param file: location of the file
        :returns: file content in bytes
        """
        with open(file, 'rb') as f:
            return f.read()

    def write_output(self, file, data):
        """
        Function to write the given content to a file in output folder

        :param file: location of the file
        :param data: Content in bytes to be written
        """
        print(f'Writing file {file} to output folder')
        with open(os.path.join(FILE_PATH, 'output', file), 'wb') as f:
            f.write(data)

    def run(self, option):
        """
        Function to perform encryption on decryption based on option provided

        :param option: 'ecrypt' for encryption, 'decrypt' for decryption
        """
        for file in os.listdir(os.path.join(FILE_PATH, 'input')):
            if file in ['key.txt', '.gitignore']:
                continue
            text = self.read_file(os.path.join(FILE_PATH, 'input', file))

            out_text = getattr(self.key, option)(text)
            self.write_output(file, out_text)

if __name__ == "__main__":
    var_map = {
        'e': 'encrypt',
        'd': 'decrypt'
    }
    print('Give "e" to encrypt files from encrypt folder, "d" decrypt files from decrypt folder')
    if sys.argv[1] not in ['e', 'd']:
        raise Exception("Invalid input!")
    e = Encryption(var_map[sys.argv[1]])
    