From ffd2d3746f8f40ceb124b675fccf28b04131120e Mon Sep 17 00:00:00 2001
From: Ilya Zhuravlev <whatever@xyz.is>
Date: Wed, 2 Dec 2020 11:37:43 -0500
Subject: [PATCH] firmware_flasher: add more progress reporting and lock UI

---
 src/main/python/firmware_flasher.py | 98 +++++++++++++++++++++--------
 src/main/python/main_window.py      | 22 +++++--
 2 files changed, 88 insertions(+), 32 deletions(-)

diff --git a/src/main/python/firmware_flasher.py b/src/main/python/firmware_flasher.py
index d478370..9019a7d 100644
--- a/src/main/python/firmware_flasher.py
+++ b/src/main/python/firmware_flasher.py
@@ -31,45 +31,56 @@ def send_retries(dev, data, retries=20):
 CHUNK = 64
 
 
-def cmd_flash(device, firmware, log):
+def cmd_flash(device, firmware, log_cb, progress_cb, complete_cb, error_cb):
+    while len(firmware) % CHUNK != 0:
+        firmware += b"\x00"
+
     # Check bootloader is correct version
     device.send(b"VC\x00")
     data = device.recv(8)
-    log("* Bootloader version: {}".format(data[0]))
+    log_cb("* Bootloader version: {}".format(data[0]))
     if data[0] != 0:
-        log("Error: Unsupported bootloader version")
-        return
+        return error_cb("Error: Unsupported bootloader version")
 
     # TODO: Check vial ID against firmware package
     device.send(b"VC\x01")
     data = device.recv(8)
-    log("* Vial ID: {}".format(data.hex()))
+    log_cb("* Vial ID: {}".format(data.hex()))
 
     # Flash
-    firmware_size = len(firmware)
-    if firmware_size % CHUNK != 0:
-        firmware_size += CHUNK - firmware_size % CHUNK
-    log(device.send(b"VC\x02" + struct.pack("<H", firmware_size // CHUNK)))
+    log_cb("Flashing...")
+    device.send(b"VC\x02" + struct.pack("<H", len(firmware) // CHUNK))
+    total = 0
     for part in chunks(firmware, CHUNK):
         if len(part) < CHUNK:
             part += b"\x00" * (CHUNK - len(part))
         if not send_retries(device, part):
-            log("Error while sending data, firmware is corrupted.")
-            return
-        log(datetime.datetime.now())
+            return error_cb("Error while sending data, firmware is corrupted")
+        total += len(part)
+        progress_cb(total / len(firmware))
 
     # Reboot
-    log("Rebooting...")
+    log_cb("Rebooting...")
     device.send(b"VC\x03")
 
+    complete_cb("Done!")
+
 
 class FirmwareFlasher(QVBoxLayout):
     log_signal = pyqtSignal(object)
+    progress_signal = pyqtSignal(object)
+    complete_signal = pyqtSignal(object)
+    error_signal = pyqtSignal(object)
 
-    def __init__(self, parent=None):
+    def __init__(self, main, parent=None):
         super().__init__(parent)
 
+        self.main = main
+
         self.log_signal.connect(self._on_log)
+        self.progress_signal.connect(self._on_progress)
+        self.complete_signal.connect(self._on_complete)
+        self.error_signal.connect(self._on_error)
 
         self.selected_firmware_path = ""
 
@@ -77,21 +88,23 @@ class FirmwareFlasher(QVBoxLayout):
         self.txt_file_selector = QLineEdit()
         self.txt_file_selector.setReadOnly(True)
         file_selector.addWidget(self.txt_file_selector)
-        btn_select_file = QToolButton()
-        btn_select_file.setText(tr("Flasher", "Select file..."))
-        btn_select_file.clicked.connect(self.on_click_select_file)
-        file_selector.addWidget(btn_select_file)
+        self.btn_select_file = QToolButton()
+        self.btn_select_file.setText(tr("Flasher", "Select file..."))
+        self.btn_select_file.clicked.connect(self.on_click_select_file)
+        file_selector.addWidget(self.btn_select_file)
         self.addLayout(file_selector)
         self.txt_logger = QPlainTextEdit()
         self.txt_logger.setReadOnly(True)
         self.txt_logger.setFont(QFontDatabase.systemFont(QFontDatabase.FixedFont))
         self.addWidget(self.txt_logger)
         progress_flash = QHBoxLayout()
-        progress_flash.addWidget(QProgressBar())
-        btn_flash = QToolButton()
-        btn_flash.setText(tr("Flasher", "Flash"))
-        btn_flash.clicked.connect(self.on_click_flash)
-        progress_flash.addWidget(btn_flash)
+        self.progress_bar = QProgressBar()
+        self.progress_bar.setRange(0, 100)
+        progress_flash.addWidget(self.progress_bar)
+        self.btn_flash = QToolButton()
+        self.btn_flash.setText(tr("Flasher", "Flash"))
+        self.btn_flash.clicked.connect(self.on_click_flash)
+        progress_flash.addWidget(self.btn_flash)
         self.addLayout(progress_flash)
 
         self.device = None
@@ -130,13 +143,46 @@ class FirmwareFlasher(QVBoxLayout):
             return
 
         self.log("Preparing to flash...")
-        threading.Thread(target=lambda: cmd_flash(self.device, firmware, self.on_log)).start()
+        self.lock_ui()
+        threading.Thread(target=lambda: cmd_flash(
+            self.device, firmware, self.on_log, self.on_progress, self.on_complete, self.on_error)).start()
 
     def on_log(self, line):
         self.log_signal.emit(line)
 
-    def _on_log(self, line):
-        self.log(line)
+    def on_progress(self, progress):
+        self.progress_signal.emit(progress)
+
+    def on_complete(self, msg):
+        self.complete_signal.emit(msg)
+
+    def on_error(self, msg):
+        self.error_signal.emit(msg)
+
+    def _on_log(self, msg):
+        self.log(msg)
+
+    def _on_progress(self, progress):
+        self.progress_bar.setValue(int(progress * 100))
+
+    def _on_complete(self, msg):
+        self.log(msg)
+        self.progress_bar.setValue(100)
+        self.unlock_ui()
+
+    def _on_error(self, msg):
+        self.log(msg)
+        self.unlock_ui()
 
     def log(self, line):
         self.txt_logger.appendPlainText("[{}] {}".format(datetime.datetime.now(), line))
+
+    def lock_ui(self):
+        self.btn_select_file.setEnabled(False)
+        self.btn_flash.setEnabled(False)
+        self.main.lock_ui()
+
+    def unlock_ui(self):
+        self.btn_select_file.setEnabled(True)
+        self.btn_flash.setEnabled(True)
+        self.main.unlock_ui()
diff --git a/src/main/python/main_window.py b/src/main/python/main_window.py
index e113da7..dabe2cb 100644
--- a/src/main/python/main_window.py
+++ b/src/main/python/main_window.py
@@ -23,17 +23,17 @@ class MainWindow(QMainWindow):
         self.combobox_devices = QComboBox()
         self.combobox_devices.currentIndexChanged.connect(self.on_device_selected)
 
-        btn_refresh_devices = QToolButton()
-        btn_refresh_devices.setToolButtonStyle(Qt.ToolButtonTextOnly)
-        btn_refresh_devices.setText(tr("MainWindow", "Refresh"))
-        btn_refresh_devices.clicked.connect(self.on_click_refresh)
+        self.btn_refresh_devices = QToolButton()
+        self.btn_refresh_devices.setToolButtonStyle(Qt.ToolButtonTextOnly)
+        self.btn_refresh_devices.setText(tr("MainWindow", "Refresh"))
+        self.btn_refresh_devices.clicked.connect(self.on_click_refresh)
 
         layout_combobox = QHBoxLayout()
         layout_combobox.addWidget(self.combobox_devices)
-        layout_combobox.addWidget(btn_refresh_devices)
+        layout_combobox.addWidget(self.btn_refresh_devices)
 
         self.layout_editor = LayoutEditor()
-        self.firmware_flasher = FirmwareFlasher()
+        self.firmware_flasher = FirmwareFlasher(self)
 
         self.tabs = QTabWidget()
         self.refresh_tabs()
@@ -139,3 +139,13 @@ class MainWindow(QMainWindow):
             self.sideload_vid = int(self.sideload_json["vendorId"], 16)
             self.sideload_pid = int(self.sideload_json["productId"], 16)
             self.on_click_refresh()
+
+    def lock_ui(self):
+        self.tabs.setEnabled(False)
+        self.combobox_devices.setEnabled(False)
+        self.btn_refresh_devices.setEnabled(False)
+
+    def unlock_ui(self):
+        self.tabs.setEnabled(True)
+        self.combobox_devices.setEnabled(True)
+        self.btn_refresh_devices.setEnabled(True)