UI: implement BenchmarkScreen's individual back handling

examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/BenchmarkScreen.kt

@@ -1,6 +1,7 @@
 package com.example.llama.revamp.ui.screens
 
 import android.llama.cpp.InferenceEngine.State
+import androidx.activity.compose.BackHandler
 import androidx.compose.foundation.background
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.Column
@@ -23,24 +24,32 @@ import androidx.compose.runtime.getValue
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.unit.dp
-import androidx.hilt.navigation.compose.hiltViewModel
 import com.example.llama.revamp.ui.components.ModelCard
+import com.example.llama.revamp.ui.components.UnloadDialogState
+import com.example.llama.revamp.ui.components.UnloadModelConfirmationDialog
 import com.example.llama.revamp.ui.theme.MonospacedTextStyle
 import com.example.llama.revamp.viewmodel.BenchmarkViewModel
 
 @Composable
 fun BenchmarkScreen(
-    onBackPressed: () -> Unit,
+    onNavigateBack: () -> Unit,
-    viewModel: BenchmarkViewModel = hiltViewModel()
+    viewModel: BenchmarkViewModel
 ) {
     val engineState by viewModel.engineState.collectAsState()
     val benchmarkResults by viewModel.benchmarkResults.collectAsState()
     val selectedModel by viewModel.selectedModel.collectAsState()
+    val unloadDialogState by viewModel.unloadDialogState.collectAsState()
 
+    // Run benchmark when entering the screen
     LaunchedEffect(selectedModel) {
         viewModel.runBenchmark()
     }
 
+    // Handle back button press
+    BackHandler {
+        viewModel.onBackPressed()
+    }
+
     Column(
         modifier = Modifier
             .fillMaxSize()
@@ -115,4 +124,35 @@ fun BenchmarkScreen(
             }
         }
     }
+
+    // Unload confirmation dialog
+    when (val state = unloadDialogState) {
+        is UnloadDialogState.Confirming -> {
+            UnloadModelConfirmationDialog(
+                onConfirm = {
+                    viewModel.onUnloadConfirmed(onNavigateBack)
+                },
+                onDismiss = {
+                    viewModel.onUnloadDismissed()
+                },
+                isUnloading = false
+            )
+        }
+        is UnloadDialogState.Unloading -> {
+            UnloadModelConfirmationDialog(
+                onConfirm = {
+                    viewModel.onUnloadConfirmed(onNavigateBack)
+                },
+                onDismiss = {
+                    viewModel.onUnloadDismissed()
+                },
+                isUnloading = true
+            )
+        }
+        is UnloadDialogState.Error -> {
+            // TODO-han.yin: TBD
+            android.util.Log.e("JOJO", "Unload error: ${state.message}")
+        }
+        else -> { /* Dialog not shown */ }
+    }
 }

examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/BenchmarkViewModel.kt

@@ -5,8 +5,11 @@ import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import com.example.llama.revamp.data.model.ModelInfo
 import com.example.llama.revamp.engine.BenchmarkService
+import com.example.llama.revamp.ui.components.UnloadDialogState
 import dagger.hilt.android.lifecycle.HiltViewModel
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.launch
 import javax.inject.Inject
 
@@ -14,10 +17,18 @@ import javax.inject.Inject
 class BenchmarkViewModel @Inject constructor(
     private val benchmarkService: BenchmarkService
 ) : ViewModel() {
-
+    /**
+     * Core states
+     */
     val engineState: StateFlow<State> = benchmarkService.engineState
-    val benchmarkResults: StateFlow<String?> = benchmarkService.benchmarkResults
     val selectedModel: StateFlow<ModelInfo?> = benchmarkService.currentSelectedModel
+    val benchmarkResults: StateFlow<String?> = benchmarkService.benchmarkResults
+
+    /**
+     * Model unloading dialog state
+     */
+    private val _unloadDialogState = MutableStateFlow<UnloadDialogState>(UnloadDialogState.Hidden)
+    val unloadDialogState: StateFlow<UnloadDialogState> = _unloadDialogState.asStateFlow()
 
     /**
      * Run benchmark with specified parameters
@@ -26,4 +37,53 @@ class BenchmarkViewModel @Inject constructor(
         viewModelScope.launch {
             benchmarkService.benchmark(pp, tg, pl, nr)
         }
+
+    /**
+     * Handle back press from both back button and top bar
+     */
+    fun onBackPressed() {
+        when (engineState.value) {
+            State.Benchmarking -> {
+                // Ignore back navigation requests during active benchmarking
+            }
+            else -> _unloadDialogState.value = UnloadDialogState.Confirming
+        }
+    }
+
+    /**
+     * Handle confirmation from unload dialog
+     */
+    fun onUnloadConfirmed(onNavigateBack: () -> Unit) {
+        viewModelScope.launch {
+            // Set unloading state to show progress
+            _unloadDialogState.value = UnloadDialogState.Unloading
+            android.util.Log.d("JOJO", "onUnloadConfirmed $ state -> Unloading")
+
+            try {
+                // Unload the model
+                benchmarkService.unloadModel()
+                android.util.Log.d("JOJO", "onUnloadConfirmed $ service unload model finished!")
+
+                // Reset state and navigate back
+                _unloadDialogState.value = UnloadDialogState.Hidden
+                android.util.Log.d("JOJO", "onUnloadConfirmed $ state -> Hidden!")
+                onNavigateBack()
+            } catch (e: Exception) {
+                // Handle error if needed
+                _unloadDialogState.value = UnloadDialogState.Hidden
+            }
+        }
+    }
+
+    /**
+     * Handle dismissal of unload dialog
+     */
+    fun onUnloadDismissed() {
+        when (_unloadDialogState.value) {
+            is UnloadDialogState.Unloading -> {
+                // Ignore dismissing requests during active benchmarking
+            }
+            else -> _unloadDialogState.value = UnloadDialogState.Hidden
+        }
+    }
 }