import sqlite3

# Function to add a new user to the database
def setupMfaSecret(user, secret):
    with sqlite3.connect('mfa.db') as conn:
        cursor = conn.cursor()
        cursor.execute('''
            INSERT INTO users (user, secret)
            VALUES (?, ?)
        ''', (user, secret))
        conn.commit()

# Function to retrieve the number corresponding to a username
def getMfaSecret(user):
    with sqlite3.connect('mfa.db') as conn:
        cursor = conn.cursor()
        cursor.execute('''
            SELECT secret FROM users
            WHERE user = ?
        ''', (user,))
        result = cursor.fetchone()
        if result:
            return result[0]
        else:
            return None

# Function to update the number for a specific username
def updateMfaSecret(user, secret):
    with sqlite3.connect('mfa.db') as conn:
        cursor = conn.cursor()
        cursor.execute('''
            UPDATE users
            SET secret = ?
            WHERE user = ?
        ''', (secret, user))
        conn.commit()

def removeMfa(user):
    with sqlite3.connect('mfa.db') as conn:
        cursor = conn.cursor()
        cursor.execute('''
            DELETE FROM users
            WHERE user = ?
        ''', (user,))
        conn.commit()

def initializeDb():
    with sqlite3.connect('mfa.db') as conn:
        cursor = conn.cursor()
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS users (
                user TEXT PRIMARY KEY,
                secret TEXT
            )
        ''')

# initializeDb()

# Example usage
# add_user('user1', 123)
# add_user('user2', 456)

# print(f"Number for 'user1': {get_number('user1')}")

# update_number('user1', 789)
# print(f"Updated number for 'user1': {get_number('user1')}")