From 618026f96cfe1d2d90925f922a37b5bf86c0d8a5 Mon Sep 17 00:00:00 2001
From: Neil Booth <kyuupichan@gmail.com>
Date: Fri, 22 Jan 2016 22:16:58 +0900
Subject: [PATCH] Installwizard: tweak hardware wallet handling

---
 gui/qt/installwizard.py | 31 +++++++++++++++++++++----------
 gui/qt/util.py          |  4 ++--
 lib/wizard.py           | 19 ++++++++++++++-----
 3 files changed, 37 insertions(+), 17 deletions(-)

diff --git a/gui/qt/installwizard.py b/gui/qt/installwizard.py
index b4a42a22..e751233e 100644
--- a/gui/qt/installwizard.py
+++ b/gui/qt/installwizard.py
@@ -253,14 +253,6 @@ class InstallWizard(WindowModalDialog, WizardBase):
         Return a a tuple (action, kind_index).  Action is 'create' or
         'restore', and kind the index of the wallet kind chosen."""
 
-        hw_wallet_help = _(
-            "NOTE regarding hardware wallets:  If you want to set up a new "
-            "or wiped device, or if you want Electrum to manage a device "
-            "you have already set up, select \"create\".  If you wish to "
-            "recover a wallet from a device seed and no longer use the "
-            "device, select \"restore\"."
-            )
-
         actions = [_("Create a new wallet"),
                    _("Restore a wallet or import keys")]
         title = _("Electrum could not find an existing wallet.")
@@ -270,13 +262,32 @@ class InstallWizard(WindowModalDialog, WizardBase):
         vbox = QVBoxLayout()
         vbox.addLayout(actions_clayout.layout())
         vbox.addLayout(wallet_clayout.layout())
-        vbox.addSpacing(10)
-        vbox.addWidget(WWLabel(hw_wallet_help))
         self.set_main_layout(vbox, title)
 
         action = ['create', 'restore'][actions_clayout.selected_index()]
         return action, wallet_clayout.selected_index()
 
+    def query_hw_wallet_choice(self, msg, action, choices):
+        actions = [_("Initialize a new or wiped device"),
+                   _("Use a device you have already set up"),
+                   _("Restore Electrum wallet from device seed words")]
+        default_action = 1 if action == 'create' else 2
+        actions_clayout = ChoicesLayout(_("What do you want to do?"), actions,
+                                        checked_index=default_action)
+        wallet_clayout = ChoicesLayout(msg, choices)
+
+        vbox = QVBoxLayout()
+        vbox.addLayout(actions_clayout.layout())
+        vbox.addLayout(wallet_clayout.layout())
+        self.set_main_layout(vbox)
+        self.next_button.setEnabled(len(choices) != 0)
+
+        if actions_clayout.selected_index() == 2:
+            action = 'restore'
+        else:
+            action = 'create'
+        return action, wallet_clayout.selected_index()
+
     def request_many(self, n, xpub_hot=None):
         vbox = QVBoxLayout()
         scroll = QScrollArea()
diff --git a/gui/qt/util.py b/gui/qt/util.py
index 4b549ec4..34759a70 100644
--- a/gui/qt/util.py
+++ b/gui/qt/util.py
@@ -245,7 +245,7 @@ def text_dialog(parent, title, label, ok_label, default=None):
         return unicode(txt.toPlainText())
 
 class ChoicesLayout(object):
-    def __init__(self, msg, choices, on_clicked=None):
+    def __init__(self, msg, choices, on_clicked=None, checked_index=0):
         vbox = QVBoxLayout()
         if len(msg) > 50:
             vbox.addWidget(WWLabel(msg))
@@ -263,7 +263,7 @@ class ChoicesLayout(object):
             vbox2.addWidget(button)
             group.addButton(button)
             group.setId(button, i)
-            if i==0:
+            if i==checked_index:
                 button.setChecked(True)
 
         if on_clicked:
diff --git a/lib/wizard.py b/lib/wizard.py
index 8b6c9bec..07d91c7d 100644
--- a/lib/wizard.py
+++ b/lib/wizard.py
@@ -80,6 +80,15 @@ class WizardBase(PrintError):
         Return the index of the choice."""
         raise NotImplementedError
 
+    def query_hw_wallet_choice(self, msg, action, choices):
+        """Asks the user which hardware wallet kind they are using.  Action is
+        'create' or 'restore' from the initial screen.  As this is
+        confusing for hardware wallets ask a new question with the
+        three possibilities Initialize ('create'), Use ('create') or
+        Restore a software-only wallet ('restore').  Return a pair
+        (action, choice)."""
+        raise NotImplementedError
+
     def show_and_verify_seed(self, seed):
         """Show the user their seed.  Ask them to re-enter it.  Return
         True on success."""
@@ -203,12 +212,12 @@ class WizardBase(PrintError):
         if kind == 'multisig':
             wallet_type = self.query_multisig(action)
         elif kind == 'hardware':
+            # The create/restore distinction is not obvious for hardware
+            # wallets; so we ask for the action again and default based
+            # on the prior choice :)
             hw_wallet_types, choices = self.plugins.hardware_wallets(action)
-            if action == 'create':
-                msg = _('Select the hardware wallet to create')
-            else:
-                msg = _('Select the hardware wallet to restore')
-            choice = self.query_choice(msg, choices)
+            msg = _('Select the type of hardware wallet: ')
+            action, choice = self.query_hw_wallet_choice(msg, action, choices)
             wallet_type = hw_wallet_types[choice]
         elif kind == 'twofactor':
             wallet_type = '2fa'