import os
import getpass
import sys
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from pathlib import Path
import base64

# TODO Add option to write and read keyfile, where there is an option to store the salt!

if_verbose = False

def verbose(string):
    if if_verbose: print(string)

# Encrypting data
def encrypt_data_neo(unencrypted_name, password, encrypted_name=None, use_key=None, keyfile_name=None, delete_original=None):

    verbose(f"original file: {unencrypted_name}")

    if not encrypted_name:
        encrypted_name="encrypted"
    
    file_chosen_name=encrypted_name
    verbose(f"new file: {encrypted_name}")

    if use_key:
        verbose("Using key")
    else:
        verbose("not using key")
    
    verbose(f"password: {password}")
    password = password.encode()
    verbose(f"password encoded: {password}")

    salt = os.urandom(16)

    verbose(f"salt: {salt}")

    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=480000,
    )

    key = base64.urlsafe_b64encode(kdf.derive(password))

    verbose(f"key: {key}")

    f = Fernet(key)

    verbose(f"f: {f}")

    try:
        with open(unencrypted_name, 'rb') as file:
            original_data = file.read()
        
        verbose(f"{unencrypted_name} has been read")

        encrypted_data = f.encrypt(original_data)

        verbose(f"data has been encrypted.")

        if not use_key:
            with open(f"{encrypted_name}", 'wb') as file:
                # Write the 16-byte salt first, then the encrypted data
                file.write(salt)
                file.write(encrypted_data)
            verbose(f"data was written to {encrypted_name} without key.")

        elif use_key:
            if not keyfile_name:
                keyfile_name = "key.p7c_key"
            
            verbose(f"keyfile_name: {keyfile_name}")

            with open(f"{keyfile_name}", 'wb') as key:
                key.write(salt)

            verbose("salt has been written.")

            with open(f"{encrypted_name}", 'wb') as file:
                file.write(encrypted_data)

            verbose("encrypted data has been written.")
        else:
            verbose("🔥🔥🔥 THIS IS AN ERROR 🔥🔥🔥")
            print(f"This is P7MJ. I am speechless. use_key: {use_key}")

        if delete_original:
            os.remove(f"{unencrypted_name}")
            verbose("Original file has been removed.")

    except Exception as e:
        verbose("🔥🔥🔥 THIS IS AN ERROR 🔥🔥🔥")
        print(f"\nAn error has occured.\n {e}\n")

def decrypt_data_neo(encrypted_name, password, unencrypted_name=None, use_key=None, keyfile_name=None, delete_original=None):

    choose_file = encrypted_name

    if not unencrypted_name:
        unencrypted_name = "unencrypted"

    password = password.encode()

    if use_key and not keyfile_name:
        keyfile_name = "key.p7c_key"

    # 2. Open the encrypted file and extract the salt + data
    if use_key:
        with open(f"{keyfile_name}", 'rb') as file:
            file_salt = file.read(16)
        with open(choose_file, 'rb') as file:
            encrypted_data = file.read()
    
    elif not use_key:
        with open(encrypted_name, 'rb') as file:
            # Read exactly 16 bytes for the salt
            file_salt = file.read(16)
            # Read everything else as the encrypted data
            encrypted_data = file.read()
    else:
        verbose("🔥🔥🔥 THIS IS AN ERROR 🔥🔥🔥 flames are burning and so is my brain")
        print("WHAT DID YOU DO TO THERE_IS_A_KEY?!")

    # 3. Re-derive the EXACT same key using that salt
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=file_salt,
        iterations=480000,
    )

    key = base64.urlsafe_b64encode(kdf.derive(password))

    verbose(f"key: {key}")

    f = Fernet(key)

    verbose(f"f: {f}")

    try:
        decrypted_data = f.decrypt(encrypted_data)
        verbose("data decrypted.")
        
        with open(f'{unencrypted_name}', 'wb') as file:
            file.write(decrypted_data)
        verbose("decrypted data written to disk.")
        if delete_original:
            os.remove(encrypted_name)
            verbose("Encrypted file removed.")
            if use_key:
                os.remove(keyfile_name)
            verbose("Key removed.")
        
    except Exception as e:
        print(f"\nError. Could not decrypt. It is likely that your password is wrong.\nError: {e}.\n")

if __name__ == "__main__":
    try:
        while True:
            inputs = input("[PCS] ")
            list_inputs = inputs.split()
            verbose(inputs)
            verbose(list_inputs)

            if not list_inputs:
                continue

            for i, param in enumerate(list_inputs):
                if param == "None" or param == "0" or param == "No" or param == "no" or param == "False":
                    list_inputs[i] = None
                if param == "True" or param == "1" or param == "Yes" or param == "yes" or param == "Exists" or param == "exists":
                    list_inputs[i] = True
            
            if list_inputs[0] == "help":
                print("Encrypt:")
                print(" encrypt unencrypted_name [encrypted_name] password [use_key] [keyfile_name] [delete_original]")
                print()
                print("Decrypt:")
                print(" decrypt encrypted_name [unencrypted_name] password [use_key] [keyfile_name] [delete_original]")
                print()
                print("Credits:")
                print(" credits")
                print()
                print("Exit:")
                print(" exit")
                print()
                print("Notes:")
                print(" [Optional]: optional parameter")
                print("\nThis program is able to encrypt absolutely ANYTHING.\nIt is recommended to name your encrypted archives after your original file.")
                print(" E.g. archive.zip should be named archive.zip.p7c_enc.")
                print(" This ensures that you remember the extension and don't end up decrypting into the wrong format.")
                print()
                print("Compatibility Notices")
                print(" For compatibility with PCS A-2-i Dividend, the keyfile must be named key.p7c_key.\n  Additionally, append .p7c_enc to all encrypted archive names.")
                print(" For compatability with PCS A-1-i Proto, follow A-2-i standards, but do not use any form of keyfile.")
            
            elif list_inputs[0] == "credits":
                print("PCS A-3-ii Terminal by P7MJ.")
            
            elif list_inputs[0] == "exit":
                sys.exit(0)

            elif list_inputs[0] == "encrypt":
                if len(list_inputs) == 7:
                    # unencrypted_name, password, encrypted_name=None, use_key=None, keyfile_name=None, delete_original=None
                    encrypt_data_neo(list_inputs[1], list_inputs[3], encrypted_name=list_inputs[2], use_key=list_inputs[4], keyfile_name=list_inputs[5], delete_original=list_inputs[6])
                else:
                    print("Invalid use of encrypt: length of parameters incorrect.")
                    print('If you wish to leave a parameter empty, type "None", "0", "No", "no", or "False".')
                    print('Type "help" for more information.')
            
            elif list_inputs[0] == "decrypt":
                if len(list_inputs) == 7:
                    # encrypted_name, password, unencrypted_name=None, use_key=None, keyfile_name=None, delete_original=None
                    decrypt_data_neo(list_inputs[1], list_inputs[3], unencrypted_name=list_inputs[2], use_key=list_inputs[4], keyfile_name=list_inputs[5], delete_original=list_inputs[6])
                else:
                    print("Invalid use of decrypt: length of parameters incorrect.")
                    print('If you wish to leave a parameter empty, type "None", "0", "No", "no", or "False".')
                    print('Type "help" for more information.')
            else:
                print("Command not found.")
    except KeyboardInterrupt:
        pass
    except Exception as e:
        print(f"{e}")