Send messages to launcher renderer to update grid (2/3)

Test: Manually tested. See bug.
Bug: 362237825
Flag: com.android.systemui.shared.new_customization_picker_ui
Change-Id: Id2046ba0900b6a2928b794e83cd92111c0a7b68d
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModel.kt
index 0fdc0b0..c3ff291 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModel.kt
@@ -49,8 +49,14 @@
     // The currently-set system grid option
     val selectedGridOption =
         interactor.selectedGridOption.filterNotNull().map { toOptionItemViewModel(it) }
-
-    private val overrideGridOptionKey = MutableStateFlow<String?>(null)
+    private val _previewingGridOptionKey = MutableStateFlow<String?>(null)
+    // If the previewing key is null, use the currently-set system grid option
+    val previewingGridOptionKey =
+        combine(selectedGridOption, _previewingGridOptionKey) {
+            currentlySetGridOption,
+            previewingGridOptionKey ->
+            previewingGridOptionKey ?: currentlySetGridOption.key.value
+        }
 
     val optionItems: Flow<List<OptionItemViewModel<GridIconViewModel>>> =
         interactor.gridOptions.filterNotNull().map { gridOptions ->
@@ -58,16 +64,20 @@
         }
 
     val onApply: Flow<(() -> Unit)?> =
-        combine(selectedGridOption, overrideGridOptionKey) {
+        combine(selectedGridOption, _previewingGridOptionKey) {
             selectedGridOption,
-            overrideGridOptionKey ->
+            previewingGridOptionKey ->
             if (
-                overrideGridOptionKey == null ||
-                    overrideGridOptionKey == selectedGridOption.key.value
+                previewingGridOptionKey == null ||
+                    previewingGridOptionKey == selectedGridOption.key.value
             ) {
                 null
             } else {
-                { viewModelScope.launch { interactor.applySelectedOption(overrideGridOptionKey) } }
+                {
+                    viewModelScope.launch {
+                        interactor.applySelectedOption(previewingGridOptionKey)
+                    }
+                }
             }
         }
 
@@ -86,7 +96,7 @@
                     )
             )
         val isSelected =
-            overrideGridOptionKey
+            _previewingGridOptionKey
                 .map {
                     if (it == null) {
                         option.isCurrent
@@ -113,7 +123,7 @@
             onClicked =
                 isSelected.map {
                     if (!it) {
-                        { overrideGridOptionKey.value = option.key }
+                        { _previewingGridOptionKey.value = option.key }
                     } else {
                         null
                     }
diff --git a/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt b/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
index 7d40b5b..fac7230 100644
--- a/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
+++ b/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
@@ -18,6 +18,7 @@
 
 import android.os.Bundle
 import android.os.Message
+import androidx.core.os.bundleOf
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.lifecycleScope
@@ -65,65 +66,86 @@
             )
         }
 
-        if (screen == Screen.LOCK_SCREEN) {
-            lifecycleOwner.lifecycleScope.launch {
-                lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
-                    launch {
-                        viewModel.selectedOption.collect {
-                            when (it) {
-                                ThemePickerLockCustomizationOption.SHORTCUTS ->
-                                    workspaceCallback.sendMessage(
-                                        MESSAGE_ID_START_CUSTOMIZING_QUICK_AFFORDANCES,
-                                        Bundle().apply {
-                                            putString(
-                                                KEY_INITIALLY_SELECTED_SLOT_ID,
-                                                SLOT_ID_BOTTOM_START,
-                                            )
-                                        }
-                                    )
-                                else ->
-                                    workspaceCallback.sendMessage(
-                                        MESSAGE_ID_DEFAULT_PREVIEW,
-                                        Bundle.EMPTY,
-                                    )
-                            }
-                        }
-                    }
-
-                    launch {
-                        viewModel.keyguardQuickAffordancePickerViewModel2.selectedSlotId.collect {
-                            workspaceCallback.sendMessage(
-                                MESSAGE_ID_SLOT_SELECTED,
-                                Bundle().apply { putString(KEY_SLOT_ID, it) },
-                            )
-                        }
-                    }
-
-                    launch {
-                        viewModel.keyguardQuickAffordancePickerViewModel2.selectedQuickAffordances
-                            .collect {
-                                it[SLOT_ID_BOTTOM_START]?.let {
-                                    workspaceCallback.sendMessage(
-                                        MESSAGE_ID_PREVIEW_QUICK_AFFORDANCE_SELECTED,
-                                        Bundle().apply {
-                                            putString(KEY_SLOT_ID, SLOT_ID_BOTTOM_START)
-                                            putString(KEY_QUICK_AFFORDANCE_ID, it)
-                                        },
-                                    )
-                                }
-                                it[SLOT_ID_BOTTOM_END]?.let {
-                                    workspaceCallback.sendMessage(
-                                        MESSAGE_ID_PREVIEW_QUICK_AFFORDANCE_SELECTED,
-                                        Bundle().apply {
-                                            putString(KEY_SLOT_ID, SLOT_ID_BOTTOM_END)
-                                            putString(KEY_QUICK_AFFORDANCE_ID, it)
-                                        },
-                                    )
+        when (screen) {
+            Screen.LOCK_SCREEN ->
+                lifecycleOwner.lifecycleScope.launch {
+                    lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+                        launch {
+                            viewModel.selectedOption.collect {
+                                when (it) {
+                                    ThemePickerLockCustomizationOption.SHORTCUTS ->
+                                        workspaceCallback.sendMessage(
+                                            MESSAGE_ID_START_CUSTOMIZING_QUICK_AFFORDANCES,
+                                            Bundle().apply {
+                                                putString(
+                                                    KEY_INITIALLY_SELECTED_SLOT_ID,
+                                                    SLOT_ID_BOTTOM_START,
+                                                )
+                                            }
+                                        )
+                                    else ->
+                                        workspaceCallback.sendMessage(
+                                            MESSAGE_ID_DEFAULT_PREVIEW,
+                                            Bundle.EMPTY,
+                                        )
                                 }
                             }
+                        }
+
+                        launch {
+                            viewModel.keyguardQuickAffordancePickerViewModel2.selectedSlotId
+                                .collect {
+                                    workspaceCallback.sendMessage(
+                                        MESSAGE_ID_SLOT_SELECTED,
+                                        Bundle().apply { putString(KEY_SLOT_ID, it) },
+                                    )
+                                }
+                        }
+
+                        launch {
+                            viewModel.keyguardQuickAffordancePickerViewModel2
+                                .selectedQuickAffordances
+                                .collect {
+                                    it[SLOT_ID_BOTTOM_START]?.let {
+                                        workspaceCallback.sendMessage(
+                                            MESSAGE_ID_PREVIEW_QUICK_AFFORDANCE_SELECTED,
+                                            Bundle().apply {
+                                                putString(KEY_SLOT_ID, SLOT_ID_BOTTOM_START)
+                                                putString(KEY_QUICK_AFFORDANCE_ID, it)
+                                            },
+                                        )
+                                    }
+                                    it[SLOT_ID_BOTTOM_END]?.let {
+                                        workspaceCallback.sendMessage(
+                                            MESSAGE_ID_PREVIEW_QUICK_AFFORDANCE_SELECTED,
+                                            Bundle().apply {
+                                                putString(KEY_SLOT_ID, SLOT_ID_BOTTOM_END)
+                                                putString(KEY_QUICK_AFFORDANCE_ID, it)
+                                            },
+                                        )
+                                    }
+                                }
+                        }
                     }
                 }
-            }
+            Screen.HOME_SCREEN ->
+                lifecycleOwner.lifecycleScope.launch {
+                    lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+                        launch {
+                            viewModel.shapeAndGridPickerViewModel.previewingGridOptionKey.collect {
+                                workspaceCallback.sendMessage(
+                                    MESSAGE_ID_UPDATE_GRID,
+                                    bundleOf(KEY_GRID_NAME to it)
+                                )
+                            }
+                        }
+                    }
+                }
         }
     }
+
+    companion object {
+        const val MESSAGE_ID_UPDATE_GRID = 7414
+        const val KEY_GRID_NAME = "grid_name"
+    }
 }