import time
from struct import pack

import electrum_ecc as ecc

from electrum.i18n import _
from electrum.util import UserCancelled, UserFacingException
from electrum.keystore import bip39_normalize_passphrase
from electrum.bip32 import BIP32Node, convert_bip32_strpath_to_intpath as parse_path
from electrum.logging import Logger
from electrum.plugin import runs_in_hwd_thread
from electrum.hw_wallet.plugin import OutdatedHwFirmwareException, HardwareClientBase

from trezorlib.client import TrezorClient, PassphraseSetting, AppManifest, get_client
from trezorlib.exceptions import TrezorFailure, Cancelled, OutdatedFirmwareError
from trezorlib.messages import WordRequestType, FailureType, ButtonRequestType, Capability, Features
from trezorlib import models
from trezorlib.thp.pairing import CodeEntry, ControllerLifecycle
import trezorlib.btc
import trezorlib.device

from trezorlib.messages import RecoveryDeviceInputMethod


MESSAGES = {
    ButtonRequestType.ConfirmOutput:
        _("Confirm the transaction output on your {} device"),
    ButtonRequestType.ResetDevice:
        _("Complete the initialization process on your {} device"),
    ButtonRequestType.ConfirmWord:
        _("Write down the seed word shown on your {}"),
    ButtonRequestType.WipeDevice:
        _("Confirm on your {} that you want to wipe it clean"),
    ButtonRequestType.ProtectCall:
        _("Confirm on your {} device the message to sign"),
    ButtonRequestType.SignTx:
        _("Confirm the total amount spent and the transaction fee on your {} device"),
    ButtonRequestType.Address:
        _("Confirm wallet address on your {} device"),
    ButtonRequestType._Deprecated_ButtonRequest_PassphraseType:
        _("Choose on your {} device where to enter your passphrase"),
    ButtonRequestType.PassphraseEntry:
        _("Please enter your passphrase on the {} device"),
    'default': _("Check your {} device to continue"),
}

# trezorlib THP does not support cross-thread cancellation.
# https://github.com/trezor/trezor-firmware/issues/7112
CANCEL_SUPPORTED = frozenset({models.T1B1, models.T2T1, models.T2B1, models.T3T1, models.T3B1})

class TrezorClientBase(HardwareClientBase, Logger):
    def __init__(self, transport, handler, plugin):
        HardwareClientBase.__init__(self, plugin=plugin)
        Logger.__init__(self)

        if plugin.is_outdated_fw_ignored():
            TrezorClient.is_outdated = lambda *args, **kwargs: False

        self._session = None
        self.device = plugin.device
        self.handler = handler

        self.transport = transport
        self.app = AppManifest(
            app_name="Electrum",
            credentials=(),
            button_callback=self.button_request,
            pin_callback=self.get_pin,
        )
        self._client = None
        # Makes sure the client is connected (on THP, a channel has been established)
        model = self.client.model
        if model.is_unknown:
            self.logger.warning("Unknown Trezor model: %s", model)

        # Pairing cannot be done during device enumeration, since UI handler is unset).

        self.msg = None
        self.creating_wallet = False

        self.in_flow = False

        self.used()

    def is_paired(self) -> bool:
        return self.client.pairing.is_paired()

    def pair_if_needed(self) -> None:
        if self.is_paired():
            return

        assert self.handler is not None, "No UI handler for pairing"

        pairing = self.client.pairing
        with self.client:
            try:
                method = CodeEntry(pairing)
                code = self.handler.get_word(_("Enter 6-digit pairing code:"))
                method.send_code(code)

                assert pairing.state is ControllerLifecycle.PAIRING_COMPLETED
                pairing.finish()
            except Exception:
                # Drop THP client (a new channel will be created later)
                self._client = None
                raise

    @property
    def session(self):
        if self._session is None:
            assert self.handler is not None, "No UI handler for session"
            self.pair_if_needed()

            # If needed, unlock the device (triggering PIN entry dialog for legacy model).
            with self.client.get_session(passphrase=PassphraseSetting.STANDARD_WALLET) as session:
                session.ensure_unlocked()

            passphrase = PassphraseSetting.STANDARD_WALLET  # (empty passphrase)
            if self.client.features.passphrase_protection:
                passphrase = self.get_passphrase(Capability.PassphraseEntry in self.client.features.capabilities)

            # Then, derive a session for this wallet (possibly with a passphrase)
            if passphrase == PassphraseSetting.STANDARD_WALLET:
                self._session = session  # reuse the session above to avoid re-derivation
                self.logger.info("Opened standard %s", self._session)
            else:
                self._session = self.client.get_session(passphrase)
                self.logger.info("Re-opened passphrase %s", self._session)

        return self._session


    def run_flow(self, message=None, creating_wallet=False):
        if self.in_flow:
            raise RuntimeError("Overlapping call to run_flow")

        self.in_flow = True
        self.msg = message
        self.creating_wallet = creating_wallet
        self.prevent_timeouts()
        return self

    def end_flow(self):
        self.in_flow = False
        self.msg = None
        self.creating_wallet = False
        self.handler.finished()
        self.used()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, e, traceback):
        self.end_flow()
        if e is not None:
            if isinstance(e, Cancelled):
                raise UserCancelled() from e
            elif isinstance(e, TrezorFailure):
                raise RuntimeError(str(e)) from e
            elif isinstance(e, OutdatedFirmwareError):
                raise OutdatedHwFirmwareException(e) from e
            else:
                return False
        return True

    @property
    def client(self) -> TrezorClient:
        if self._client is None:
            # Connect to the device, without pairing (on THP)
            self._client = get_client(self.app, self.transport)
        return self._client

    @property
    def features(self):
        assert self.is_paired(), "No features"
        return self.client.features

    def label(self):
        if not self.is_paired():
            return None
        return self.features.label

    def get_soft_device_id(self):
        if not self.is_paired():
            return None
        return self.features.device_id

    def is_initialized(self) -> bool | None:
        if not self.is_paired():
            return None  # Pairing will be done later

        return bool(self.features.initialized)

    def is_pairable(self):
        if not self.is_paired():
            return True
        return not self.features.bootloader_mode

    @runs_in_hwd_thread
    def has_usable_connection_with_device(self):
        if self.in_flow:
            return True

        try:
            self.client.ping(message="")
        except Exception as e:
            self.logger.exception("No connection: %s", e)
            return False
        return True

    def used(self):
        self.last_operation = time.time()

    def prevent_timeouts(self):
        self.last_operation = float('inf')

    @runs_in_hwd_thread
    def timeout(self, cutoff):
        '''Time out the client if the last operation was before cutoff.'''
        if self.last_operation < cutoff:
            self.logger.info("timed out")
            self.clear_session()

    def i4b(self, x):
        return pack('>I', x)

    @runs_in_hwd_thread
    def get_xpub(self, bip32_path, xtype, creating=False):
        address_n = parse_path(bip32_path)
        with self.run_flow(creating_wallet=creating):
            node = trezorlib.btc.get_public_node(self.session, address_n).node
        return BIP32Node(xtype=xtype,
                         eckey=ecc.ECPubkey(node.public_key),
                         chaincode=node.chain_code,
                         depth=node.depth,
                         fingerprint=self.i4b(node.fingerprint),
                         child_number=self.i4b(node.child_num)).to_xpub()

    @runs_in_hwd_thread
    def toggle_passphrase(self):
        if self.features.passphrase_protection:
            msg = _("Confirm on your {} device to disable passphrases")
        else:
            msg = _("Confirm on your {} device to enable passphrases")
        enabled = not self.features.passphrase_protection
        with self.run_flow(msg):
            trezorlib.device.apply_settings(self.session, use_passphrase=enabled)

    @runs_in_hwd_thread
    def change_label(self, label):
        with self.run_flow(_("Confirm the new label on your {} device")):
            trezorlib.device.apply_settings(self.session, label=label)

    @runs_in_hwd_thread
    def change_homescreen(self, homescreen):
        with self.run_flow(_("Confirm on your {} device to change your home screen")):
            trezorlib.device.apply_settings(self.session, homescreen=homescreen)

    @runs_in_hwd_thread
    def set_pin(self, remove):
        if remove:
            msg = _("Confirm on your {} device to disable PIN protection")
        elif self.features.pin_protection:
            msg = _("Confirm on your {} device to change your PIN")
        else:
            msg = _("Confirm on your {} device to set a PIN")
        with self.run_flow(msg):
            trezorlib.device.change_pin(self.session, remove)

    @runs_in_hwd_thread
    def clear_session(self):
        '''Clear the session to force pin (and passphrase if enabled)
        re-entry.  Does not leak exceptions.'''
        self.logger.info(f"clear session: {self}")
        self.prevent_timeouts()
        try:
            self.close()
        except BaseException as e:
            # If the device was removed it has the same effect...
            self.logger.info(f"clear_session: ignoring error {e}")

    @runs_in_hwd_thread
    def close(self):
        '''Called when Our wallet was closed or the device removed.'''
        self.logger.info("locking: %s", self.client)
        self.client.lock()
        self.logger.info("closing: %s", self._session)
        if self._session is not None:
            self._session.close()

    @runs_in_hwd_thread
    def is_uptodate(self):
        if self.client.is_outdated():
            return False
        return self.client.version >= self.plugin.minimum_firmware

    def get_trezor_model(self):
        return self.client.model.name

    def device_model_name(self):
        model = self.get_trezor_model()
        if model == '1':
            return "Trezor One"
        elif model == 'T':
            return "Trezor T"
        elif model == "Safe 3":
            return "Trezor Safe 3"
        elif model == "Safe 5":
            return "Trezor Safe 5"
        elif model == "Safe 7":
            return "Trezor Safe 7"
        return None

    @runs_in_hwd_thread
    def show_address(self, address_str, script_type, multisig=None):
        coin_name = self.plugin.get_coin_name()
        address_n = parse_path(address_str)
        with self.run_flow():
            return trezorlib.btc.get_address(
                self.session,
                coin_name,
                address_n,
                show_display=True,
                script_type=script_type,
                multisig=multisig)

    @runs_in_hwd_thread
    def sign_message(self, address_str, message, *, script_type):
        coin_name = self.plugin.get_coin_name()
        address_n = parse_path(address_str)
        with self.run_flow():
            return trezorlib.btc.sign_message(
                self.session,
                coin_name,
                address_n,
                message,
                script_type=script_type,
                no_script_type=True)

    @runs_in_hwd_thread
    def recover_device(self, recovery_type, *args, **kwargs):
        input_callback = self.mnemonic_callback(recovery_type)
        with self.run_flow(), self.client.get_session(None) as seedless_session:
            return trezorlib.device.recover(
                seedless_session,
                *args,
                input_callback=input_callback,
                type=recovery_type,
                **kwargs)

    # ========= Unmodified trezorlib methods =========

    @runs_in_hwd_thread
    def sign_tx(self, *args, **kwargs):
        with self.run_flow():
            return trezorlib.btc.sign_tx(self.session, *args, **kwargs)

    @runs_in_hwd_thread
    def get_ownership_id(self, *args, **kwargs):
        with self.run_flow():
            return trezorlib.btc.get_ownership_id(self.session, *args, **kwargs)

    @runs_in_hwd_thread
    def get_ownership_proof(self, *args, **kwargs):
        with self.run_flow():
            return trezorlib.btc.get_ownership_proof(self.session, *args, **kwargs)

    @runs_in_hwd_thread
    def reset_device(self, *args, **kwargs):
        with self.run_flow(), self.client.get_session(None) as seedless_session:
            return trezorlib.device.reset(seedless_session, *args, **kwargs)

    @runs_in_hwd_thread
    def wipe_device(self, *args, **kwargs):
        with self.run_flow(), self.client.get_session(None) as seedless_session:
            return trezorlib.device.wipe(seedless_session, *args, **kwargs)

    # ========= UI methods ==========

    def button_request(self, br):
        message = self.msg or MESSAGES.get(br.code) or MESSAGES['default']
        on_cancel = self.client.cancel if self.client.model in CANCEL_SUPPORTED else None
        self.handler.show_message(message.format(self.device), on_cancel)

    def get_pin(self, code=None):
        show_strength = True
        if code == 2:
            msg = _("Enter a new PIN for your {}:")
        elif code == 3:
            msg = (_("Re-enter the new PIN for your {}.\n\n"
                     "NOTE: the positions of the numbers have changed!"))
        else:
            msg = _("Enter your current {} PIN:")
            show_strength = False
        pin = self.handler.get_pin(msg.format(self.device), show_strength=show_strength)
        if not pin:
            raise Cancelled
        # check PIN length. Depends on model and firmware version
        # https://github.com/trezor/trezor-firmware/issues/1167
        limit = 9
        if self.get_trezor_model() == "1":
            if (1, 10, 0) <= self.client.version:
                limit = 50
        else:
            if (2, 4, 0) <= self.client.version:
                limit = 50
        if len(pin) > limit:
            self.handler.show_error(_('The PIN cannot be longer than {} characters.').format(limit))
            raise Cancelled
        return pin

    def get_passphrase(self, available_on_device):
        if self.creating_wallet:
            msg = _("Enter a passphrase to generate this wallet.  Each time "
                    "you use this wallet your {} will prompt you for the "
                    "passphrase.  If you forget the passphrase you cannot "
                    "access the bitcoins in the wallet.").format(self.device)
        else:
            msg = _("Enter the passphrase to unlock this wallet:")

        self.handler.passphrase_on_device = available_on_device
        passphrase = self.handler.get_passphrase(msg, self.creating_wallet)
        if passphrase is PassphraseSetting.ON_DEVICE:
            return passphrase
        if passphrase is None:
            raise Cancelled
        passphrase = bip39_normalize_passphrase(passphrase)
        length = len(passphrase)
        if length > 50:
            self.handler.show_error(_("Too long passphrase ({} > 50 chars).").format(length))
            raise Cancelled
        return passphrase

    def _matrix_char(self, matrix_type):
        num = 9 if matrix_type == WordRequestType.Matrix9 else 6
        char = self.handler.get_matrix(num)
        if char == 'x':
            raise Cancelled
        return char

    def mnemonic_callback(self, recovery_type):
        if recovery_type is None:
            return None

        if recovery_type == RecoveryDeviceInputMethod.Matrix:
            return self._matrix_char

        step = 0
        def word_callback(_ignored):
            nonlocal step
            step += 1
            msg = _("Step {}/24.  Enter seed word as explained on your {}:").format(step, self.device)
            word = self.handler.get_word(msg)
            if not word:
                raise Cancelled
            return word
        return word_callback
