import json
import os
from getpass import getpass

from tonsdk.contract.wallet import Wallets
from tonsdk.crypto import mnemonic_to_wallet_key
from tonsdk.utils import Address

from _core import unpack_wallet, encrypt_data, decrypt_data


def mnemonic_new(subwallet_id: int=0):
    while True:
        mnemonic, public_key, private_key, _ = Wallets.create('v3r2', 0)
        wallet = unpack_wallet({
            'mnemonic': mnemonic,
            'public_key': public_key.hex(),
            'secret_key': private_key.hex(),
            'subwallet_id': subwallet_id,
            'custom_address': None,
            'is_testnet': False,
        })
        ugly = False
        for wallet_address in [
            wallet.address.to_string(1, 1, 0),
            wallet.address.to_string(1, 1, 1)
        ]:
            if ('_' in wallet_address or '-' in wallet_address):
                ugly = True
                
        if not ugly:
            return mnemonic


def get_or_create_wallet_settings():
    wallet_settings_file = os.environ.get('ENCRYPTED_WALLET_FILE', '.saved_wallet')
    
    # Check if the .saved_wallet file exists
    if os.path.exists(wallet_settings_file):
        # If the file exists, read and decrypt its contents
        with open(wallet_settings_file, 'rb') as file:
            _content = file.read()
            
        hashpart = _content[:32].hex()
        _content = _content[32:]
        print(f"=== Saved wallet: {Address('0:' + hashpart).to_string(1, 1, 0)}")
        
        password = getpass("Enter the password to decrypt the wallet: ")
        try:
            wallet_settings = decrypt_data(_content, password.encode("UTF-8"))
            wallet_settings = json.loads(wallet_settings.decode())
        except (ValueError, KeyError):
            print("Incorrect password or corrupted data. Please try again.")
            return get_or_create_wallet_settings()

        print(f"=== Mainnet: {not wallet_settings.get('is_testnet', True)}")
        return wallet_settings
    else:
        # If the file doesn't exist, prompt the user for input
        try:
            subwallet_id = int(input("Enter the subwallet_id (default is 0): ") or 0)
            assert subwallet_id >= 0, "Subwallet ID must be a positive integer"
            
            mnemonic = input("Enter the wallet mnemonic (24 words, separated by spaces): ")
            if mnemonic == 'new':
                mnemonic = ' '.join(mnemonic_new(subwallet_id=subwallet_id))
                
            assert len(mnemonic.split()) == 24, "Mnemonic must be 24 words long"
            public_key, secret_key = mnemonic_to_wallet_key(mnemonic.split())
            
            custom_address = input("Enter a custom address (if available): ")
            custom_address = Address(custom_address) if custom_address else None
            password = getpass("Enter a password to protect the wallet: ")
            assert len(password) > 3, "Password must be at least 4 characters long"
            is_testnet = input("Is this a testnet wallet? (yes/NO): ").lower() == 'yes'
        except KeyboardInterrupt:
            os._exit(1)
        except BaseException as e:
            print(f"Error: {e}")
            print("Please try again")
            return get_or_create_wallet_settings()

        # Create a dictionary with the user input
        wallet_settings = {
            'mnemonic': mnemonic,
            'public_key': public_key.hex(),
            'secret_key': secret_key.hex(),
            'subwallet_id': subwallet_id,
            'custom_address': custom_address.to_string(1, 1, 1) if custom_address else None,
            'is_testnet': is_testnet,
        }
        print(f"=== Mainnet: {not is_testnet}")
        wallet_address = unpack_wallet(wallet_settings).address.to_string(0, 0, 0)

        # Encrypt and save the data to the file
        encrypted_wallet = encrypt_data(json.dumps(wallet_settings).encode("UTF-8"), password.encode("UTF-8"))
        with open(wallet_settings_file, 'wb') as file:
            file.write(bytes.fromhex(wallet_address.split(':')[1]) + encrypted_wallet)

        # Recursively call the function to read data from the file
        return get_or_create_wallet_settings()