From 04fb390bee30bd0009c1304ccd7211658cd4ec13 Mon Sep 17 00:00:00 2001
From: Arsen Musayelyan <moussaelianarsen@gmail.com>
Date: Sat, 6 Nov 2021 19:06:17 -0700
Subject: [PATCH] Add reminder to validate firmware to itctl and itgui

---
 cmd/itctl/firmware/upgrade.go |  8 ++++++++
 cmd/itgui/upgrade.go          | 20 ++++++++++++++++++--
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/cmd/itctl/firmware/upgrade.go b/cmd/itctl/firmware/upgrade.go
index 7df69dd..37489e5 100644
--- a/cmd/itctl/firmware/upgrade.go
+++ b/cmd/itctl/firmware/upgrade.go
@@ -19,6 +19,9 @@
 package firmware
 
 import (
+	"fmt"
+	"time"
+
 	"github.com/cheggaaa/pb/v3"
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
@@ -38,6 +41,8 @@ var upgradeCmd = &cobra.Command{
 	Short:   "Upgrade InfiniTime firmware using files or archive",
 	Aliases: []string{"upg"},
 	Run: func(cmd *cobra.Command, args []string) {
+		start := time.Now()
+
 		client := viper.Get("client").(*api.Client)
 
 		var upgType api.UpgradeType
@@ -79,6 +84,9 @@ var upgradeCmd = &cobra.Command{
 		}
 		// Finish progress bar
 		bar.Finish()
+
+		fmt.Printf("Transferred %d B in %s.\n", bar.Total(), time.Since(start))
+		fmt.Println("Remember to validate the new firmware in the InfiniTime settings.")
 	},
 }
 
diff --git a/cmd/itgui/upgrade.go b/cmd/itgui/upgrade.go
index 30a89c1..f9b31db 100644
--- a/cmd/itgui/upgrade.go
+++ b/cmd/itgui/upgrade.go
@@ -134,8 +134,6 @@ func upgradeTab(parent fyne.Window, client *api.Client) *fyne.Container {
 
 		// Show progress dialog
 		progressDlg.Show()
-		// Hide progress dialog after completion
-		defer progressDlg.Hide()
 
 		for event := range progress {
 			// Set label text to received / total B
@@ -150,6 +148,24 @@ func upgradeTab(parent fyne.Window, client *api.Client) *fyne.Container {
 				break
 			}
 		}
+
+		// Hide progress dialog after completion
+		progressDlg.Hide()
+
+		// Reset screen to default
+		upgradeTypeSelect.SetSelectedIndex(0)
+		firmwareBtn.SetText("Select firmware (.bin)")
+		initPktBtn.SetText("Select init packet (.dat)")
+		archiveBtn.SetText("Select archive (.zip)")
+		firmwarePath = ""
+		initPktPath = ""
+		archivePath = ""
+
+		dialog.NewInformation(
+			"Upgrade Complete",
+			"The firmware was transferred successfully.\nRemember to validate the firmware in InfiniTime settings.",
+			parent,
+		).Show()
 	})
 
 	// Return container containing all elements