From a8dc825aefaa5abec1869fe220a6caeadc66d622 Mon Sep 17 00:00:00 2001
From: Han Yin <han.yin@arm.com>
Date: Sun, 20 Apr 2025 19:51:32 -0700
Subject: [PATCH] UI: handle back press on Model Selection screen

---
 .../java/com/example/llama/revamp/MainActivity.kt  |  3 +++
 .../revamp/ui/screens/ModelSelectionScreen.kt      | 14 ++++++++++----
 .../revamp/viewmodel/ModelSelectionViewModel.kt    |  9 +++++++++
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/examples/llama.android/app/src/main/java/com/example/llama/revamp/MainActivity.kt b/examples/llama.android/app/src/main/java/com/example/llama/revamp/MainActivity.kt
index e1532039dd..c9b429071f 100644
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/MainActivity.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/MainActivity.kt
@@ -293,6 +293,9 @@ fun AppContent(
                 // Model Selection Screen
                 composable(AppDestinations.MODEL_SELECTION_ROUTE) {
                     ModelSelectionScreen(
+                        onNavigateBack = {
+                            navigationActions.navigateUp()
+                        },
                         onModelConfirmed = { modelInfo ->
                             navigationActions.navigateToModelLoading()
                         },
diff --git a/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/ModelSelectionScreen.kt b/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/ModelSelectionScreen.kt
index 87f807a5a2..875e772d50 100644
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/ModelSelectionScreen.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/ModelSelectionScreen.kt
@@ -1,5 +1,6 @@
 package com.example.llama.revamp.ui.screens
 
+import androidx.activity.compose.BackHandler
 import androidx.compose.animation.AnimatedVisibility
 import androidx.compose.animation.fadeIn
 import androidx.compose.animation.fadeOut
@@ -47,9 +48,9 @@ fun ModelSelectionScreen(
     val models by viewModel.availableModels.collectAsState()
     val preselectedModel by viewModel.preselectedModel.collectAsState()
 
-    val handleModelSelection = { model: ModelInfo ->
-        viewModel.confirmSelectedModel(model)
-        onModelConfirmed(model)
+    // Handle back button press
+    BackHandler(preselectedModel != null) {
+        viewModel.onBackPressed()
     }
 
     Box(modifier = Modifier.fillMaxSize()) {
@@ -89,7 +90,12 @@ fun ModelSelectionScreen(
             exit = scaleOut() + fadeOut()
         ) {
             FloatingActionButton(
-                onClick = { preselectedModel?.let { handleModelSelection(it) } },
+                onClick = {
+                    preselectedModel?.let {
+                        viewModel.confirmSelectedModel(it)
+                        onModelConfirmed(it)
+                    }
+                },
                 containerColor = MaterialTheme.colorScheme.primary
             ) {
                 Icon(
diff --git a/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/ModelSelectionViewModel.kt b/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/ModelSelectionViewModel.kt
index 304422f110..352db8502e 100644
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/ModelSelectionViewModel.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/ModelSelectionViewModel.kt
@@ -55,6 +55,15 @@ class ModelSelectionViewModel @Inject constructor(
         _preselectedModel.value = null
     }
 
+    /**
+     * Handle back press from both back button and top bar
+     */
+    fun onBackPressed() {
+        if (_preselectedModel.value != null) {
+            resetSelection()
+        }
+    }
+
     companion object {
         private val TAG = ModelSelectionViewModel::class.java.simpleName