Merge "Color Picker Optimistic Update" into tm-qpr-dev am: 5189d47f43

Original change: https://googleplex-android-review.googlesource.com/c/platform/packages/apps/ThemePicker/+/21424417

Change-Id: I14e9f3c20e092f87b2390071becdf149e8138ccf
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt b/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
index c375574..0e65577 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
@@ -25,7 +25,15 @@
  * system color.
  */
 interface ColorPickerRepository {
+    /**
+     * The newly selected color option for overwriting the current active option during an
+     * optimistic update, the value is null when no overwriting is needed
+     */
+    val activeColorOption: Flow<ColorOptionModel?>
+
+    /** List of wallpaper and preset color options on the device, categorized by Color Type */
     val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>>
 
+    /** Selects a color option with optimistic update */
     fun select(colorOptionModel: ColorOptionModel)
 }
diff --git a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
index 4ce5ed9..d6d5060 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
@@ -29,7 +29,9 @@
 import com.android.customization.picker.color.shared.model.ColorType
 import com.android.wallpaper.model.WallpaperColorsViewModel
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.suspendCancellableCoroutine
@@ -48,12 +50,17 @@
     private val colorManager: ColorCustomizationManager =
         ColorCustomizationManager.getInstance(context, OverlayManagerCompat(context))
 
-    /** List of wallpaper and preset color options on the device, categorized by Color Type */
+    private val _activeColorOption = MutableStateFlow<ColorOptionModel?>(null)
+    override val activeColorOption: StateFlow<ColorOptionModel?> = _activeColorOption.asStateFlow()
+
     override val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>> =
-        combine(homeWallpaperColors, lockWallpaperColors) { homeColors, lockColors ->
-                homeColors to lockColors
+        combine(activeColorOption, homeWallpaperColors, lockWallpaperColors) {
+                activeOption,
+                homeColors,
+                lockColors ->
+                Triple(activeOption, homeColors, lockColors)
             }
-            .map { (homeColors, lockColors) ->
+            .map { (activeOption, homeColors, lockColors) ->
                 suspendCancellableCoroutine { continuation ->
                     colorManager.setWallpaperColors(homeColors, lockColors)
                     colorManager.fetchOptions(
@@ -66,8 +73,9 @@
                                 options?.forEach { option ->
                                     when (option) {
                                         is ColorSeedOption ->
-                                            wallpaperColorOptions.add(option.toModel())
-                                        is ColorBundle -> presetColorOptions.add(option.toModel())
+                                            wallpaperColorOptions.add(option.toModel(activeOption))
+                                        is ColorBundle ->
+                                            presetColorOptions.add(option.toModel(activeOption))
                                     }
                                 }
                                 continuation.resumeWith(
@@ -95,23 +103,32 @@
             }
 
     override fun select(colorOptionModel: ColorOptionModel) {
+        _activeColorOption.value = colorOptionModel
         val colorOption: ColorOption = colorOptionModel.colorOption
         colorManager.apply(
             colorOption,
             object : CustomizationManager.Callback {
-                override fun onSuccess() = Unit
+                override fun onSuccess() {
+                    _activeColorOption.value = null
+                }
 
                 override fun onError(throwable: Throwable?) {
+                    _activeColorOption.value = null
                     Log.w(TAG, "Apply theme with error", throwable)
                 }
             }
         )
     }
 
-    private fun ColorOption.toModel(): ColorOptionModel {
+    private fun ColorOption.toModel(activeColorOption: ColorOptionModel?): ColorOptionModel {
         return ColorOptionModel(
             colorOption = this,
-            isSelected = isActive(colorManager),
+            isSelected =
+                if (activeColorOption != null) {
+                    isEquivalent(activeColorOption.colorOption)
+                } else {
+                    isActive(colorManager)
+                },
         )
     }
 
diff --git a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt b/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
index 331d635..d2a25bc 100644
--- a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
@@ -27,6 +27,9 @@
 import kotlinx.coroutines.flow.asStateFlow
 
 class FakeColorPickerRepository(context: Context) : ColorPickerRepository {
+    override val activeColorOption: StateFlow<ColorOptionModel?> =
+        MutableStateFlow<ColorOptionModel?>(null)
+
     private val colorSeedOption0: ColorSeedOption =
         ColorSeedOption.Builder()
             .setLightColors(