4 files changed, 84 insertions, 4 deletions

diff --git a/screen_ui.cpp b/screen_ui.cpp
index dc596314c..522aa6b23 100644
--- a/screen_ui.cpp
+++ b/screen_ui.cpp
@@ -51,6 +51,7 @@ static double now() {
 }
 
 ScreenRecoveryUI::ScreenRecoveryUI() :
+    currentIcon(NONE),
     installingFrame(0),
     locale(nullptr),
     rtl_locale(false),
@@ -75,8 +76,7 @@ ScreenRecoveryUI::ScreenRecoveryUI() :
     animation_fps(-1),
     installing_frames(-1),
     stage(-1),
-    max_stage(-1),
-    currentIcon(NONE) {
+    max_stage(-1) {
 
     for (int i = 0; i < 5; i++) {
         backgroundIcon[i] = nullptr;
diff --git a/screen_ui.h b/screen_ui.h
index 386deac2d..08a5f44a9 100644
--- a/screen_ui.h
+++ b/screen_ui.h
@@ -68,6 +68,7 @@ class ScreenRecoveryUI : public RecoveryUI {
     void SetColor(UIElement e);
 
   private:
+    Icon currentIcon;
     int installingFrame;
     const char* locale;
     bool rtl_locale;
@@ -138,8 +139,6 @@ class ScreenRecoveryUI : public RecoveryUI {
     void LoadBitmap(const char* filename, GRSurface** surface);
     void LoadBitmapArray(const char* filename, int* frames, int* fps, GRSurface*** surface);
     void LoadLocalizedBitmap(const char* filename, GRSurface** surface);
-  protected:
-    Icon currentIcon;
 };
 
 #endif  // RECOVERY_UI_H
diff --git a/wear_ui.cpp b/wear_ui.cpp
index ef48b788b..65bcd8494 100644
--- a/wear_ui.cpp
+++ b/wear_ui.cpp
@@ -59,6 +59,7 @@ WearRecoveryUI::WearRecoveryUI() :
     intro_frames(22),
     loop_frames(60),
     animation_fps(30),
+    currentIcon(NONE),
     intro_done(false),
     current_frame(0),
     rtl_locale(false),
@@ -365,6 +366,57 @@ void WearRecoveryUI::Init()
     RecoveryUI::Init();
 }
 
+void WearRecoveryUI::SetBackground(Icon icon)
+{
+    pthread_mutex_lock(&updateMutex);
+    currentIcon = icon;
+    update_screen_locked();
+    pthread_mutex_unlock(&updateMutex);
+}
+
+void WearRecoveryUI::SetProgressType(ProgressType type)
+{
+    pthread_mutex_lock(&updateMutex);
+    if (progressBarType != type) {
+        progressBarType = type;
+    }
+    progressScopeStart = 0;
+    progressScopeSize = 0;
+    progress = 0;
+    update_screen_locked();
+    pthread_mutex_unlock(&updateMutex);
+}
+
+void WearRecoveryUI::ShowProgress(float portion, float seconds)
+{
+    pthread_mutex_lock(&updateMutex);
+    progressBarType = DETERMINATE;
+    progressScopeStart += progressScopeSize;
+    progressScopeSize = portion;
+    progressScopeTime = now();
+    progressScopeDuration = seconds;
+    progress = 0;
+    update_screen_locked();
+    pthread_mutex_unlock(&updateMutex);
+}
+
+void WearRecoveryUI::SetProgress(float fraction)
+{
+    pthread_mutex_lock(&updateMutex);
+    if (fraction < 0.0) fraction = 0.0;
+    if (fraction > 1.0) fraction = 1.0;
+    if (progressBarType == DETERMINATE && fraction > progress) {
+        // Skip updates that aren't visibly different.
+        int width = progress_bar_width;
+        float scale = width * progressScopeSize;
+        if ((int) (progress * scale) != (int) (fraction * scale)) {
+            progress = fraction;
+            update_screen_locked();
+        }
+    }
+    pthread_mutex_unlock(&updateMutex);
+}
+
 void WearRecoveryUI::SetStage(int current, int max)
 {
 }
@@ -447,6 +499,16 @@ int WearRecoveryUI::SelectMenu(int sel) {
     return sel;
 }
 
+void WearRecoveryUI::EndMenu() {
+    int i;
+    pthread_mutex_lock(&updateMutex);
+    if (show_menu > 0 && text_rows > 0 && text_cols > 0) {
+        show_menu = 0;
+        update_screen_locked();
+    }
+    pthread_mutex_unlock(&updateMutex);
+}
+
 bool WearRecoveryUI::IsTextVisible()
 {
     pthread_mutex_lock(&updateMutex);
@@ -477,6 +539,13 @@ void WearRecoveryUI::ShowText(bool visible)
     pthread_mutex_unlock(&updateMutex);
 }
 
+void WearRecoveryUI::Redraw()
+{
+    pthread_mutex_lock(&updateMutex);
+    update_screen_locked();
+    pthread_mutex_unlock(&updateMutex);
+}
+
 void WearRecoveryUI::ShowFile(FILE* fp) {
     std::vector<long> offsets;
     offsets.push_back(ftell(fp));
diff --git a/wear_ui.h b/wear_ui.h
index 63a257244..35ea51660 100644
--- a/wear_ui.h
+++ b/wear_ui.h
@@ -24,6 +24,13 @@ class WearRecoveryUI : public ScreenRecoveryUI {
     WearRecoveryUI();
 
     void Init();
+    // overall recovery state ("background image")
+    void SetBackground(Icon icon);
+
+    // progress indicator
+    void SetProgressType(ProgressType type);
+    void ShowProgress(float portion, float seconds);
+    void SetProgress(float fraction);
 
     void SetStage(int current, int max);
 
@@ -42,6 +49,9 @@ class WearRecoveryUI : public ScreenRecoveryUI {
     void StartMenu(const char* const * headers, const char* const * items,
                            int initial_selection);
     int SelectMenu(int sel);
+    void EndMenu();
+
+    void Redraw();
 
     enum UIElement { HEADER, MENU, MENU_SEL_BG, MENU_SEL_FG, LOG, TEXT_FILL };
     virtual void SetColor(UIElement e);
@@ -68,6 +78,8 @@ class WearRecoveryUI : public ScreenRecoveryUI {
     int animation_fps;
 
   private:
+    Icon currentIcon;
+
     bool intro_done;
 
     int current_frame;