try:
    import base64
    import pickle
    from pysqlcipher3 import dbapi2 as sc
    from math import log2
    from password_strength import PasswordStats
    from getpass import getpass as gp
    from secrets import token_urlsafe
    from random import randint as rint, SystemRandom as sr
    from atexit import register
    from os import urandom, path, remove
except ModuleNotFoundError:
    print("You have not installed the required modules. Follow these steps to do so:\n\n1. Open the terminal (Linux/MacOS) or command prompt (Windows).\n2. Navigate to this directory and then to the files directory.\n3. Type 'pip install -r dependencies.txt'.\n4. Restart the program.\n\nIf you have followed all the steps correctly, keyvault will work on the next start.")
    exit()

def database_enc():
    conn = sc.connect("database/keyvault.db")
    cursor = conn.cursor()
    if path.isfile("database/keyvault.db"):
        for _ in range(3):
            try:
                password = gp(prompt = "Enter master password: ")
                conn.execute(f"PRAGMA key = {password}")
                conn.execute('''
                CREATE TABLE IF NOT EXISTS data (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                service TEXT,
                username TEXT,
                email TEXT,
                password TEXT,
                website TEXT,
                category TEXT,
                notes TEXT,
                totp TEXT )
                ''')
                passCorrect = True
                break
            except sc.DatabaseError:
                print("Incorrect password.\n")
                passCorrect = False
        
        if not passCorrect:
            print("You have entered a wrong password three times. Please restart the program to try again.")
            exit()
        
    else:
        print("You have not setup a master password yet. Please set one below.\n")
        while True:
            mp, mp2 = gp(prompt = "Enter a secure master password (hidden for privacy!): "), gp(prompt = "Please enter it again: ")
            if mp == mp2:
                if len(mp) < 8:
                    print("\nThe master password you have set is too weak. Please set another one.")
                else:
                    break
            else:
                print("Both of the passwords are different. Please enter the same password.")
        conn.execute(f"PRAGMA key = {mp}")
        conn.commit()

def gen():
    userpass = input("Type 'u' to generate usernames or 'p' for passwords: ").lower()
    if userpass == 'u':
        with open("files/generation/wordlist.txt", "r") as f:
            words = f.readlines()
        word1, word2 = rint(0, 8874), rint(0, 8874)
        username = f"{words[word1][0:-1]}{words[word2][0:-1]}{rint(0, 100000)}"
        print(username)
    elif userpass == 'p':
        while True:
            length = int(input("Enter password length (above 8 only): "))
            if length <= 7:
                print("The password is too short. Please enter it again.")
            else:
                break
        pool = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789`~@#$%^&*()-_=+]}[{\"';:.>,</?"
        strength = ''.join(sr().choice(pool) for i in range(length))
        print(strength)
    else:
        print("Incorrect input.")

def strength():
    password = gp(prompt = "\nEnter the password to check its strength (hidden for privacy!): ")
    if password in common_passwords:
        print("Your password is in the list of 1000 of the most common passwords. Change your password immediately.")
    else:
        poolsize = 0
        for l in password:
            if l.islower():
                poolsize += 26
            elif l.isupper():
                poolsize += 26
            elif l.isdigit():
                poolsize += 10
            else:
                poolsize += 32
        entropy = log2(poolsize**len(password))
        passstrength = PasswordStats(strength).strength()
        if passstrength < 0.33:
            print(f"[PASSWORD STRENGTH]: {passstrength}\n[PASSWORD ENTROPY]: {entropy} bits\nYour password is extremely weak. Please change it.\n")
        elif passstrength >= 0.33 and passstrength < 0.66:
            print(f"[PASSWORD STRENGTH]: {passstrength}\n[PASSWORD ENTROPY]: {entropy} bits\nYour password is of fair strength. Please consider changing it to a stronger one.\n")
        elif passstrength >= 0.66 and passstrength < 0.80:
            print(f"[PASSWORD STRENGTH]: {passstrength}\n[PASSWORD ENTROPY]: {entropy} bits\nYour password is strong.\n")
        elif passstrength >= 0.80 and passstrength < 0.9:
            print(f"[PASSWORD STRENGTH]: {passstrength}\n[PASSWORD ENTROPY]: {entropy} bits\nYour password is incredibly strong.\n")
        else:
            print(f"[PASSWORD STRENGTH]: {passstrength}\n[PASSWORD ENTROPY]: {entropy} bits\nYour password is practically uncrackable.\n")

global db
global conn
global cursor

database_enc()

with open("files/strength/common-passwords.txt", "r") as f:
    common_passwords = f.read()

print("\nkeyvault initialized.")

print("keyvault is ready to use! Type 'help' for a list of commands.\n")

while True:
    command = input("> ").lower()
    if command == 'help':
        print("\nUsage:\n\nls - list entires\nrm - remove entry\ngen - generate a password or username\nstrength - check password strength\nedit - edit an entry\nshow/view - view an entry\nfind/search - open search wizard\nhelp - display this message\nversion - print current version\n")
    elif command == 'version':
        print("keyvault: v1.0.0")
    elif command == 'gen':
        gen()
    elif command == 'strength':
        strength()
    elif command == 'exit':
        print("Thank for you using keyvault!")
        exit()