try:
    import base64
    import pickle
    from tabulate import tabulate
    from pysqlcipher3 import dbapi2 as sc
    from pyotp import TOTP
    from time import time
    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
    from os import path
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():
    global conn
    global cursor
    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")

def add():
    service = username = email = password = website = category = notes = totp = "  "
    print("For any of these fields, don't type anything to make it an empty field.")
    service = input("Enter service: ")
    username = input("Enter username: ")
    email = input("Enter email: ")
    password = gp(prompt = "Enter password (hidden for privacy): ")
    website = input("Enter website: ")
    category = input("Enter category: ")
    notes = input("Enter notes: ")
    totp = input("Enter TOTP secret: ")
    data = [service, username, email, password, website, category, notes, totp]
    for _ in range(len(data)):
        if data[_] == "":
            data[_] = None
    cursor.execute("INSERT INTO data(service, username, email, password, website, category, notes, totp) VALUES(?, ?, ?, ?, ?, ?, ?, ?)", tuple(data))
    conn.commit()

def ls(cmd="SELECT * FROM data"):
    cursor.execute(cmd)
    data = list(cursor.fetchall())
    if data == []:
        print("The database is empty.")
    else:
        for _ in range(len(data)):
            data[_] = list(data[_])
            if data[_][-1] != None:
                data[_].append(TOTP(str(data[_][-1])).now())
            else:
                data[_].append(None)

        head = ["ID", "Service", "Username", "Email", "Password", "Website", "Category", "Notes", "TOTP Secret", "TOTP Code"]
        print(tabulate(data, headers = head, tablefmt = "grid"))
    return data

def rm():
    ls()
    delid = int(input("Enter the ID of the record you want to delete: "))
    cursor.execute("DELETE FROM data WHERE id=?", (delid,))
    conn.commit()

def find():
    print("The following search criteria are available: id, service, username, email, website, category")
    criterian = input("Enter your criterian here: ").lower()
    if criterian not in ['id', 'service', 'username', 'email', 'website', 'category']:
        print("You have selected an invalid criterian.")
    else:
        srchquery = input("Enter your search query: ")
        sqlquery = f"SELECT * FROM data WHERE {criterian}='{srchquery}'"
        ls(sqlquery)

def edit():
    print("The following edit criteria are available: id, service, username, email, password, website, category, notes, totp")
    criterian = input("Enter your criterian here: ").lower()
    if criterian not in ['id', 'service', 'username', 'email', 'password', 'website', 'category', 'notes', 'totp']:
        print("You have selected an invalid criterian.")
    else:
        idsearch = input("Enter the ID of the entry to edit: ")
        data = ls(f"SELECT * FROM data WHERE id={idsearch}")
        if data == []:
            print("This entry does not exist.")
        else:
            confirmation = input("Is this the correct entry? [y/n] > ").lower()
            if confirmation == 'y':
                newvalue = input("Enter the new value: ")
                cursor.execute(f"UPDATE data SET {criterian}='{newvalue}' WHERE id={idsearch}")
            else:
                print("You have cancelled the entry edit process.")

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\nadd - add an entry\nls - list entires\nrm - remove entry\ngen - generate a password or username\nstrength - check password strength\nedit - edit an entry\nfind/search - open search wizard\nhelp - display this message\nversion - print current version\n")
    elif command == 'search' or command == 'find':
        find()
    elif command == 'ls':
        ls()
    elif command == 'add':
        add()
    elif command == 'rm':
        rm()
    elif command == 'edit':
        edit()
    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()