#
# BitBox02 Electrum plugin code.
#

import hid
import hashlib
from typing import TYPE_CHECKING, Dict, Tuple, Optional, List, Any

from electrum import bip32, constants
from electrum.i18n import _
from electrum.keystore import Hardware_KeyStore, Xpub
from electrum.transaction import PartialTransaction
from electrum.wallet import Standard_Wallet, Multisig_Wallet, Deterministic_Wallet
from electrum.util import bh2u, UserFacingException
from electrum.base_wizard import ScriptTypeNotSupported, BaseWizard
from electrum.logging import get_logger
from electrum.crypto import hmac_oneshot
from electrum.plugin import Device, DeviceInfo
from electrum.simple_config import SimpleConfig
from electrum.json_db import StoredDict
from electrum.storage import get_derivation_used_for_hw_device_encryption

import electrum.bitcoin as bitcoin
import electrum.ecc as ecc

from ..hw_wallet import HW_PluginBase, HardwareClientBase
from ..hw_wallet.plugin import LibraryFoundButUnusable


try:
    from bitbox02 import bitbox02
    from bitbox02 import util
    from bitbox02.communication import (
        devices,
        HARDENED,
        u2fhid,
        bitbox_api_protocol,
    )
    requirements_ok = True
except ImportError:
    requirements_ok = False


_logger = get_logger(__name__)


class BitBox02Client(HardwareClientBase):
    # handler is a BitBox02_Handler, importing it would lead to a circular dependency
    def __init__(self, handler: Any, device: Device, config: SimpleConfig):
        self.bitbox02_device = None
        self.handler = handler
        self.device_descriptor = device
        self.config = config
        self.bitbox_hid_info = None
        if self.config.get("bitbox02") is None:
            bitbox02_config: dict = {
                "remote_static_noise_keys": [],
                "noise_privkey": None,
            }
            self.config.set_key("bitbox02", bitbox02_config)

        bitboxes = devices.get_any_bitbox02s()
        for bitbox in bitboxes:
            if (
                bitbox["path"] == self.device_descriptor.path
                and bitbox["interface_number"]
                == self.device_descriptor.interface_number
            ):
                self.bitbox_hid_info = bitbox
        if self.bitbox_hid_info is None:
            raise Exception("No BitBox02 detected")

    def label(self) -> str:
        return "BitBox02"

    def is_initialized(self) -> bool:
        return True

    def close(self):
        try:
            self.bitbox02_device.close()
        except:
            pass

    def has_usable_connection_with_device(self) -> bool:
        if self.bitbox_hid_info is None:
                return False
        return True

    def pairing_dialog(self, wizard: bool = True):
        def pairing_step(code):
            msg = "Please compare and confirm the pairing code on your BitBox02:\n"
            choice = [code]
            if wizard == True:
                return self.handler.win.query_choice(msg, choice)
            self.handler.pairing_code_dialog(code)

        def exists_remote_static_pubkey(pubkey: bytes) -> bool:
            bitbox02_config = self.config.get("bitbox02")
            noise_keys = bitbox02_config.get("remote_static_noise_keys")
            if noise_keys is not None:
                if pubkey.hex() in [noise_key for noise_key in noise_keys]:
                    return True
            return False

        def set_remote_static_pubkey(pubkey: bytes) -> None:
            if not exists_remote_static_pubkey(pubkey):
                bitbox02_config = self.config.get("bitbox02")
                if bitbox02_config.get("remote_static_noise_keys") is not None:
                    bitbox02_config["remote_static_noise_keys"].append(pubkey.hex())
                else:
                    bitbox02_config["remote_static_noise_keys"] = [pubkey.hex()]
                self.config.set_key("bitbox02", bitbox02_config)

        def get_noise_privkey() -> Optional[bytes]:
            bitbox02_config = self.config.get("bitbox02")
            privkey = bitbox02_config.get("noise_privkey")
            if privkey is not None:
                return bytes.fromhex(privkey)
            return None

        def set_noise_privkey(privkey: bytes) -> None:
            bitbox02_config = self.config.get("bitbox02")
            bitbox02_config["noise_privkey"] = privkey.hex()
            self.config.set_key("bitbox02", bitbox02_config)

        def attestation_warning() -> None:
            self.handler.attestation_failed_warning(
                "The BitBox02 attestation failed.\nTry reconnecting the BitBox02.\nWarning: The device might not be genuine, if the\n problem persists please contact Shift support."
            )

        class NoiseConfig(bitbox_api_protocol.BitBoxNoiseConfig):
            """NoiseConfig extends BitBoxNoiseConfig"""

            def show_pairing(self, code: str) -> bool:
                choice = [code]
                try:
                    reply = pairing_step(code)
                except:
                    # Close the hid device on exception
                    hid_device.close()
                    raise
                return True

            def attestation_check(self, result: bool) -> None:
                if not result:
                    attestation_warning()

            def contains_device_static_pubkey(self, pubkey: bytes) -> bool:
                return exists_remote_static_pubkey(pubkey)

            def add_device_static_pubkey(self, pubkey: bytes) -> None:
                return set_remote_static_pubkey(pubkey)

            def get_app_static_privkey(self) -> Optional[bytes]:
                return get_noise_privkey()

            def set_app_static_privkey(self, privkey: bytes) -> None:
                return set_noise_privkey(privkey)

        if self.bitbox02_device is None:
            hid_device = hid.device()
            hid_device.open_path(self.bitbox_hid_info["path"])

            self.bitbox02_device = bitbox02.BitBox02(
                transport=u2fhid.U2FHid(hid_device),
                device_info=self.bitbox_hid_info,
                noise_config=NoiseConfig(),
            )

        if not self.bitbox02_device.device_info()["initialized"]:
            raise Exception(
                "Please initialize the BitBox02 using the BitBox app first before using the BitBox02 in electrum"
            )

    def check_device_firmware_version(self) -> bool:
        if self.bitbox02_device is None:
            raise Exception(
                "Need to setup communication first before attempting any BitBox02 calls"
            )
        return self.bitbox02_device.check_firmware_version()

    def coin_network_from_electrum_network(self) -> int:
        if constants.net.TESTNET:
            return bitbox02.btc.TBTC
        return bitbox02.btc.BTC

    def get_password_for_storage_encryption(self) -> str:
        derivation = get_derivation_used_for_hw_device_encryption()
        derivation_list = bip32.convert_bip32_path_to_list_of_uint32(derivation)
        xpub = self.bitbox02_device.electrum_encryption_key(derivation_list)
        node = bip32.BIP32Node.from_xkey(xpub, net = constants.BitcoinMainnet()).subkey_at_public_derivation(())
        return node.eckey.get_public_key_bytes(compressed=True).hex()

    def get_xpub(self, bip32_path: str, xtype: str, display: bool = False) -> str:
        if self.bitbox02_device is None:
            self.pairing_dialog(wizard=False)

        if self.bitbox02_device is None:
            raise Exception(
                "Need to setup communication first before attempting any BitBox02 calls"
            )

        if not self.bitbox02_device.device_info()["initialized"]:
            raise UserFacingException(
                "Please initialize the BitBox02 using the BitBox app first before using the BitBox02 in electrum"
            )

        xpub_keypath = bip32.convert_bip32_path_to_list_of_uint32(bip32_path)
        coin_network = self.coin_network_from_electrum_network()

        if xtype == "p2wpkh":
            if coin_network == bitbox02.btc.BTC:
                out_type = bitbox02.btc.BTCPubRequest.ZPUB
            else:
                out_type = bitbox02.btc.BTCPubRequest.VPUB
        elif xtype == "p2wpkh-p2sh":
            if coin_network == bitbox02.btc.BTC:
                out_type = bitbox02.btc.BTCPubRequest.YPUB
            else:
                out_type = bitbox02.btc.BTCPubRequest.UPUB
        elif xtype == "p2wsh":
            if coin_network == bitbox02.btc.BTC:
                out_type = bitbox02.btc.BTCPubRequest.CAPITAL_ZPUB
            else:
                out_type = bitbox02.btc.BTCPubRequest.CAPITAL_VPUB
        # The other legacy types are not supported
        else:
            raise Exception("invalid xtype:{}".format(xtype))

        return self.bitbox02_device.btc_xpub(
            keypath=xpub_keypath,
            xpub_type=out_type,
            coin=coin_network,
            display=display,
        )

    def request_root_fingerprint_from_device(self) -> str:
        if self.bitbox02_device is None:
            raise Exception(
                "Need to setup communication first before attempting any BitBox02 calls"
            )

        return self.bitbox02_device.root_fingerprint().hex()

    def is_pairable(self) -> bool:
        if self.bitbox_hid_info is None:
            return False
        return True

    def btc_multisig_config(
        self, coin, bip32_path: List[int], wallet: Multisig_Wallet
    ):
        """
        Set and get a multisig config with the current device and some other arbitrary xpubs.
        Registers it on the device if not already registered.
        """

        if self.bitbox02_device is None:
            raise Exception(
                "Need to setup communication first before attempting any BitBox02 calls"
            )

        account_keypath = bip32_path[:4]
        xpubs = wallet.get_master_public_keys()
        our_xpub = self.get_xpub(
            bip32.convert_bip32_intpath_to_strpath(account_keypath), "p2wsh"
        )

        multisig_config = bitbox02.btc.BTCScriptConfig(
            multisig=bitbox02.btc.BTCScriptConfig.Multisig(
                threshold=wallet.m,
                xpubs=[util.parse_xpub(xpub) for xpub in xpubs],
                our_xpub_index=xpubs.index(our_xpub),
            )
        )

        is_registered = self.bitbox02_device.btc_is_script_config_registered(
            coin, multisig_config, account_keypath
        )
        if not is_registered:
            name = self.handler.name_multisig_account()
            try:
                self.bitbox02_device.btc_register_script_config(
                    coin=coin,
                    script_config=multisig_config,
                    keypath=account_keypath,
                    name=name,
                )
            except bitbox02.DuplicateEntryException:
                raise
            except:
                raise UserFacingException("Failed to register multisig\naccount configuration on BitBox02")
        return multisig_config

    def show_address(
        self, bip32_path: str, address_type: str, wallet: Deterministic_Wallet
    ) -> str:

        if self.bitbox02_device is None:
            raise Exception(
                "Need to setup communication first before attempting any BitBox02 calls"
            )

        address_keypath = bip32.convert_bip32_path_to_list_of_uint32(bip32_path)
        coin_network = self.coin_network_from_electrum_network()

        if address_type == "p2wpkh":
            script_config = bitbox02.btc.BTCScriptConfig(
                simple_type=bitbox02.btc.BTCScriptConfig.P2WPKH
            )
        elif address_type == "p2wpkh-p2sh":
            script_config = bitbox02.btc.BTCScriptConfig(
                simple_type=bitbox02.btc.BTCScriptConfig.P2WPKH_P2SH
            )
        elif address_type == "p2wsh":
            if type(wallet) is Multisig_Wallet:
                script_config = self.btc_multisig_config(
                    coin_network, address_keypath, wallet
                )
            else:
                raise Exception("Can only use p2wsh with multisig wallets")
        else:
            raise Exception(
                "invalid address xtype: {} is not supported by the BitBox02".format(
                    address_type
                )
            )

        return self.bitbox02_device.btc_address(
            keypath=address_keypath,
            coin=coin_network,
            script_config=script_config,
            display=True,
        )

    def sign_transaction(
        self,
        keystore: Hardware_KeyStore,
        tx: PartialTransaction,
        wallet: Deterministic_Wallet,
    ):
        if tx.is_complete():
            return

        if self.bitbox02_device is None:
            raise Exception(
                "Need to setup communication first before attempting any BitBox02 calls"
            )

        coin = bitbox02.btc.BTC
        if constants.net.TESTNET:
            coin = bitbox02.btc.TBTC

        tx_script_type = None

        # Build BTCInputType list
        inputs = []
        for txin in tx.inputs():
            _, full_path = keystore.find_my_pubkey_in_txinout(txin)

            if full_path is None:
                raise Exception(
                    "A wallet owned pubkey was not found in the transaction input to be signed"
                )

            inputs.append(
                {
                    "prev_out_hash": txin.prevout.txid[::-1],
                    "prev_out_index": txin.prevout.out_idx,
                    "prev_out_value": txin.value_sats(),
                    "sequence": txin.nsequence,
                    "keypath": full_path,
                }
            )

            if tx_script_type == None:
                tx_script_type = txin.script_type
            elif tx_script_type != txin.script_type:
                raise Exception("Cannot mix different input script types")

        if tx_script_type == "p2wpkh":
            tx_script_type = bitbox02.btc.BTCScriptConfig(
                simple_type=bitbox02.btc.BTCScriptConfig.P2WPKH
            )
        elif tx_script_type == "p2wpkh-p2sh":
            tx_script_type = bitbox02.btc.BTCScriptConfig(
                simple_type=bitbox02.btc.BTCScriptConfig.P2WPKH_P2SH
            )
        elif tx_script_type == "p2wsh":
            if type(wallet) is Multisig_Wallet:
                tx_script_type = self.btc_multisig_config(coin, full_path, wallet)
            else:
                raise Exception("Can only use p2wsh with multisig wallets")
        else:
            raise UserFacingException(
                "invalid input script type: {} is not supported by the BitBox02".format(
                    tx_script_type
                )
            )

        # Build BTCOutputType list
        outputs = []
        for txout in tx.outputs():
            assert txout.address
            # check for change
            if txout.is_change:
                _, change_pubkey_path = keystore.find_my_pubkey_in_txinout(txout)
                outputs.append(
                    bitbox02.BTCOutputInternal(
                        keypath=change_pubkey_path, value=txout.value,
                    )
                )
            else:
                addrtype, pubkey_hash = bitcoin.address_to_hash(txout.address)
                if addrtype == bitcoin.WIF_SCRIPT_TYPES["p2pkh"]:
                    output_type = bitbox02.btc.P2PKH
                elif addrtype == bitcoin.WIF_SCRIPT_TYPES["p2sh"]:
                    output_type = bitbox02.btc.P2SH
                elif addrtype == bitcoin.WIF_SCRIPT_TYPES["p2wpkh"]:
                    output_type = bitbox02.btc.P2WPKH
                elif addrtype == bitcoin.WIF_SCRIPT_TYPES["p2wsh"]:
                    output_type = bitbox02.btc.P2WSH
                else:
                    raise UserFacingException(
                        "Received unsupported output type during transaction signing: {} is not supported by the BitBox02".format(
                            addrtype
                        )
                    )
                outputs.append(
                    bitbox02.BTCOutputExternal(
                        output_type=output_type,
                        output_hash=pubkey_hash,
                        value=txout.value,
                    )
                )

        if type(wallet) is Standard_Wallet:
            keypath_account = full_path[:3]
        elif type(wallet) is Multisig_Wallet:
            keypath_account = full_path[:4]
        else:
            raise Exception(
                "BitBox02 does not support this wallet type: {}".format(type(wallet))
            )

        sigs = self.bitbox02_device.btc_sign(
            coin,
            tx_script_type,
            keypath_account=keypath_account,
            inputs=inputs,
            outputs=outputs,
            locktime=tx.locktime,
            version=tx.version,
        )

        # Fill signatures
        if len(sigs) != len(tx.inputs()):
            raise Exception("Incorrect number of inputs signed.")  # Should never occur
        signatures = [bh2u(ecc.der_sig_from_sig_string(x[1])) + "01" for x in sigs]
        tx.update_signatures(signatures)


class BitBox02_KeyStore(Hardware_KeyStore):
    hw_type = "bitbox02"
    device = "BitBox02"
    plugin: "BitBox02Plugin"

    def __init__(self, d: StoredDict):
        super().__init__(d)
        self.force_watching_only = False
        self.ux_busy = False

    def get_client(self):
        return self.plugin.get_client(self)

    def give_error(self, message: Exception, clear_client: bool = False):
        self.logger.info(message)
        if not self.ux_busy:
            self.handler.show_error(message)
        else:
            self.ux_busy = False
        if clear_client:
            self.client = None
        raise UserFacingException(message)

    def decrypt_message(self, pubkey, message, password):
        raise UserFacingException(
            _(
                "Message encryption, decryption and signing are currently not supported for {}"
            ).format(self.device)
        )

    def sign_message(self, sequence, message, password):
        raise UserFacingException(
            _(
                "Message encryption, decryption and signing are currently not supported for {}"
            ).format(self.device)
        )

    def sign_transaction(self, tx: PartialTransaction, password: str):
        if tx.is_complete():
            return
        client = self.get_client()

        try:
            try:
                self.handler.show_message("Authorize Transaction...")
                client.sign_transaction(self, tx, self.handler.win.wallet)

            finally:
                self.handler.finished()

        except Exception as e:
            self.logger.exception("")
            self.give_error(e, True)
            return

    def show_address(
        self, sequence: Tuple[int, int], txin_type: str, wallet: Deterministic_Wallet
    ):
        client = self.get_client()
        address_path = "{}/{}/{}".format(
            self.get_derivation_prefix(), sequence[0], sequence[1]
        )
        try:
            try:
                self.handler.show_message(_("Showing address ..."))
                dev_addr = client.show_address(address_path, txin_type, wallet)
            finally:
                self.handler.finished()
        except Exception as e:
            self.logger.exception("")
            self.handler.show_error(e)

class BitBox02Plugin(HW_PluginBase):
    keystore_class = BitBox02_KeyStore

    DEVICE_IDS = [(0x03EB, 0x2403)]

    SUPPORTED_XTYPES = ("p2wpkh-p2sh", "p2wpkh", "p2wsh")

    def __init__(self, parent: HW_PluginBase, config: SimpleConfig, name: str):
        super().__init__(parent, config, name)

        self.libraries_available = self.check_libraries_available()
        if not self.libraries_available:
            return
        self.device_manager().register_devices(self.DEVICE_IDS)

    def get_library_version(self):
        try:
            from bitbox02 import bitbox02
            version = bitbox02.__version__
        except:
            version = "unknown"
        if requirements_ok:
            return version
        else:
            raise ImportError()


    # handler is a BitBox02_Handler
    def create_client(self, device: Device, handler: Any) -> BitBox02Client:
        if not handler:
            self.handler = handler
        return BitBox02Client(handler, device, self.config)

    def setup_device(
        self, device_info: DeviceInfo, wizard: BaseWizard, purpose: int
    ):
        devmgr = self.device_manager()
        device_id = device_info.device.id_
        client = devmgr.client_by_id(device_id)
        if client is None:
            raise UserFacingException(
                _("Failed to create a client for this device.")
                + "\n"
                + _("Make sure it is in the correct state.")
            )
        client.handler = self.create_handler(wizard)
        if client.bitbox02_device is None:
            client.pairing_dialog()

    def get_xpub(
        self, device_id: bytes, derivation: str, xtype: str, wizard: BaseWizard
    ):
        if xtype not in self.SUPPORTED_XTYPES:
            raise ScriptTypeNotSupported(
                _("This type of script is not supported with {}.").format(self.device)
            )
        devmgr = self.device_manager()
        client = devmgr.client_by_id(device_id)
        if client.bitbox02_device is None:
            client.handler = self.create_handler(wizard)
            client.pairing_dialog()
        return client.get_xpub(derivation, xtype)

    def get_client(self, keystore: BitBox02_KeyStore, force_pair: bool = True):
        devmgr = self.device_manager()
        handler = keystore.handler
        with devmgr.hid_lock:
            client = devmgr.client_for_keystore(self, handler, keystore, force_pair)

        return client

    def show_address(
        self,
        wallet: Deterministic_Wallet,
        address: str,
        keystore: BitBox02_KeyStore = None,
    ):
        if keystore is None:
            keystore = wallet.get_keystore()
        if not self.show_address_helper(wallet, address, keystore):
            return

        txin_type = wallet.get_txin_type(address)
        sequence = wallet.get_address_index(address)
        keystore.show_address(sequence, txin_type, wallet)