#!/usr/bin/env python3
"""
Bitcoin Wallet Generation Script

This script generates a Bitcoin wallet and stores it in a KeePassXC database file (KDBX format).
The wallet includes a BIP-0039 mnemonic seed phrase, Bitcoin address, and a QR code representation
of the seed phrase. The database file is password protected to ensure the security of the wallet.

Usage: ./generate_btc_wallet.py
"""

import base64
import gc
from getpass import getpass
import qrcode
from mnemonic import Mnemonic
from bitcoinlib.keys import HDKey
from pykeepass import create_database

# Generate a BIP-0039 mnemonic seed phrase
mnemonic = Mnemonic("english")
SEED_PHRASE = mnemonic.generate(strength=128)

# Derive the HDKey from the seed phrase
hd_key = HDKey.from_passphrase(SEED_PHRASE)

# Derive the Bitcoin address from the HDKey
child_key = hd_key.subkey_for_path("m/0/0")
address = child_key.address()

# Create a QR code instance
qr = qrcode.QRCode(
    version=1,
    error_correction=qrcode.constants.ERROR_CORRECT_H,
    box_size=10,
    border=4,
)

# Add the data to the QR code
qr.add_data(SEED_PHRASE)

# Generate the QR code image
qr.make(fit=True)
QR_IMAGE = qr.make_image(fill_color="black", back_color="white")

# Convert the QR code image to base64
qr_base64 = base64.b64encode(QR_IMAGE.tobytes()).decode()

# Prompt for custom name for the wallet
wallet_name = input(
    "Whould you like to name this wallet? (empty for using the address as name): "
).strip()

# Create the database filename with the wallet number
if wallet_name == "":
    wallet_name = address
else:
    wallet_name = wallet_name[:100].strip()

db_filename = f"{wallet_name}.kdbx"

# Prompt the user for the passphrase
passphrase = getpass("Enter passphrase for the KeePassXC database: ")

try:
    # Create a KeePassXC database file
    db = create_database(db_filename, password=passphrase)

    # Create an entry in the root group
    entry = db.add_entry(db.root_group, wallet_name, username=address, password=SEED_PHRASE)

    # Add the QR code as a note (base64 encoded)
    entry.notes = qr_base64

    # Save the database
    db.save()
except Exception as e:
    print("Error while creating keepassxc database. Disk full? Readonly?")
finally:
    del SEED_PHRASE, address, hd_key, passphrase, QR_IMAGE, qr_base64, mnemonic
    gc.collect()

print("---")
print("Bitcoin address was successfully created. You can find at: " + db_filename)
print("---")