Integrate Color Picker into Preview Loading Animation (2/3)

Integrate color picker into preview loading animation so that changes in
the color section also starts the loading animation.

Bug: 274443705
Test: Manually verified
Change-Id: I618e7a99152e5e44cc6eff30748adea6065cde61
diff --git a/src/com/android/customization/module/DefaultCustomizationSections.java b/src/com/android/customization/module/DefaultCustomizationSections.java
index 04d331e..4cc82eb 100644
--- a/src/com/android/customization/module/DefaultCustomizationSections.java
+++ b/src/com/android/customization/module/DefaultCustomizationSections.java
@@ -19,6 +19,7 @@
 import com.android.customization.model.themedicon.domain.interactor.ThemedIconSnapshotRestorer;
 import com.android.customization.picker.clock.ui.view.ClockViewFactory;
 import com.android.customization.picker.clock.ui.viewmodel.ClockCarouselViewModel;
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor;
 import com.android.customization.picker.color.ui.section.ColorSectionController2;
 import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel;
 import com.android.customization.picker.notifications.ui.section.NotificationSectionController;
@@ -61,6 +62,7 @@
     private final DarkModeSnapshotRestorer mDarkModeSnapshotRestorer;
     private final ThemedIconSnapshotRestorer mThemedIconSnapshotRestorer;
     private final ThemedIconInteractor mThemedIconInteractor;
+    private final ColorPickerInteractor mColorPickerInteractor;
 
     public DefaultCustomizationSections(
             ColorPickerViewModel.Factory colorPickerViewModelFactory,
@@ -73,7 +75,8 @@
             ClockViewFactory clockViewFactory,
             DarkModeSnapshotRestorer darkModeSnapshotRestorer,
             ThemedIconSnapshotRestorer themedIconSnapshotRestorer,
-            ThemedIconInteractor themedIconInteractor) {
+            ThemedIconInteractor themedIconInteractor,
+            ColorPickerInteractor colorPickerInteractor) {
         mColorPickerViewModelFactory = colorPickerViewModelFactory;
         mKeyguardQuickAffordancePickerInteractor = keyguardQuickAffordancePickerInteractor;
         mKeyguardQuickAffordancePickerViewModelFactory =
@@ -85,6 +88,7 @@
         mDarkModeSnapshotRestorer = darkModeSnapshotRestorer;
         mThemedIconSnapshotRestorer = themedIconSnapshotRestorer;
         mThemedIconInteractor = themedIconInteractor;
+        mColorPickerInteractor = colorPickerInteractor;
     }
 
     @Override
@@ -121,6 +125,7 @@
                         sectionNavigationController,
                         wallpaperInteractor,
                         mThemedIconInteractor,
+                        mColorPickerInteractor,
                         wallpaperManager,
                         isTwoPaneAndSmallWidth,
                         savedInstanceState)
@@ -134,6 +139,7 @@
                                 wallpaperPreviewNavigator,
                                 wallpaperInteractor,
                                 mThemedIconInteractor,
+                                mColorPickerInteractor,
                                 wallpaperManager,
                                 isTwoPaneAndSmallWidth,
                                 savedInstanceState));
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index a611d7a..653c403 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -156,6 +156,7 @@
                     getDarkModeSnapshotRestorer(activity),
                     getThemedIconSnapshotRestorer(activity),
                     getThemedIconInteractor(),
+                    getColorPickerInteractor(activity, getWallpaperColorsViewModel()),
                 )
                 .also { customizationSections = it }
     }
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 7cf9fd0..fccaa65 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
@@ -19,12 +19,15 @@
 import com.android.customization.picker.color.shared.model.ColorOptionModel
 import com.android.customization.picker.color.shared.model.ColorType
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.StateFlow
 
 /**
  * Abstracts access to application state related to functionality for selecting, picking, or setting
  * system color.
  */
 interface ColorPickerRepository {
+    /** Whether the system color is in the process of being updated */
+    val isApplyingSystemColor: StateFlow<Boolean>
 
     /** List of wallpaper and preset color options on the device, categorized by Color Type */
     val colorOptions: Flow<Map<ColorType, List<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 41ef3a5..ce5c0c4 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
@@ -27,7 +27,9 @@
 import com.android.wallpaper.model.WallpaperColorsModel
 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
@@ -43,9 +45,17 @@
         wallpaperColorsViewModel.homeWallpaperColors
     private val lockWallpaperColors: StateFlow<WallpaperColorsModel?> =
         wallpaperColorsViewModel.lockWallpaperColors
+    private var selectedColorOption: MutableStateFlow<ColorOptionModel> =
+        MutableStateFlow(getCurrentColorOption())
+
+    private val _isApplyingSystemColor = MutableStateFlow(false)
+    override val isApplyingSystemColor = _isApplyingSystemColor.asStateFlow()
 
     override val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>> =
-        combine(homeWallpaperColors, lockWallpaperColors) { homeColors, lockColors ->
+        combine(homeWallpaperColors, lockWallpaperColors, selectedColorOption) {
+                homeColors,
+                lockColors,
+                _ ->
                 homeColors to lockColors
             }
             .map { (homeColors, lockColors) ->
@@ -109,17 +119,21 @@
                 }
             }
 
-    override suspend fun select(colorOptionModel: ColorOptionModel) =
+    override suspend fun select(colorOptionModel: ColorOptionModel) {
+        _isApplyingSystemColor.value = true
         suspendCancellableCoroutine { continuation ->
             colorManager.apply(
                 colorOptionModel.colorOption,
                 object : CustomizationManager.Callback {
                     override fun onSuccess() {
+                        _isApplyingSystemColor.value = false
+                        selectedColorOption.value = colorOptionModel
                         continuation.resumeWith(Result.success(Unit))
                     }
 
                     override fun onError(throwable: Throwable?) {
                         Log.w(TAG, "Apply theme with error", throwable)
+                        _isApplyingSystemColor.value = false
                         continuation.resumeWith(
                             Result.failure(throwable ?: Throwable("Error loading theme bundles"))
                         )
@@ -127,6 +141,7 @@
                 }
             )
         }
+    }
 
     override fun getCurrentColorOption(): ColorOptionModel {
         val overlays = colorManager.currentOverlays
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 714129d..bb2ef9d 100644
--- a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
@@ -29,6 +29,9 @@
 
 class FakeColorPickerRepository(private val context: Context) : ColorPickerRepository {
 
+    private val _isApplyingSystemColor = MutableStateFlow(false)
+    override val isApplyingSystemColor = _isApplyingSystemColor.asStateFlow()
+
     private lateinit var selectedColorOption: ColorOptionModel
 
     private val _colorOptions =
diff --git a/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt b/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt
index 8c7a4b7..d3b2eba 100644
--- a/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt
+++ b/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt
@@ -26,6 +26,8 @@
     private val repository: ColorPickerRepository,
     private val snapshotRestorer: Provider<ColorPickerSnapshotRestorer>,
 ) {
+    val isApplyingSystemColor = repository.isApplyingSystemColor
+
     /**
      * The newly selected color option for overwriting the current active option during an
      * optimistic update, the value is set to null when update fails
diff --git a/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt b/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt
index 7cf3550..9b1e324 100644
--- a/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt
+++ b/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt
@@ -40,6 +40,7 @@
 import com.android.customization.picker.clock.ui.view.ClockCarouselView
 import com.android.customization.picker.clock.ui.view.ClockViewFactory
 import com.android.customization.picker.clock.ui.viewmodel.ClockCarouselViewModel
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
 import com.android.wallpaper.R
 import com.android.wallpaper.model.CustomizationSectionController.CustomizationSectionNavigationController
 import com.android.wallpaper.model.WallpaperColorsViewModel
@@ -71,6 +72,7 @@
     private val navigationController: CustomizationSectionNavigationController,
     wallpaperInteractor: WallpaperInteractor,
     themedIconInteractor: ThemedIconInteractor,
+    colorPickerInteractor: ColorPickerInteractor,
     wallpaperManager: WallpaperManager,
     private val isTwoPaneAndSmallWidth: Boolean,
     savedInstanceState: Bundle?,
@@ -85,6 +87,7 @@
         wallpaperPreviewNavigator,
         wallpaperInteractor,
         themedIconInteractor,
+        colorPickerInteractor,
         wallpaperManager,
         isTwoPaneAndSmallWidth,
         savedInstanceState,
diff --git a/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt b/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
index ba48b53..b2c5815 100644
--- a/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
+++ b/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
@@ -23,6 +23,7 @@
 import android.os.Bundle
 import androidx.lifecycle.LifecycleOwner
 import com.android.customization.model.themedicon.domain.interactor.ThemedIconInteractor
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
 import com.android.customization.picker.preview.ui.viewmodel.PreviewWithThemeViewModel
 import com.android.wallpaper.R
 import com.android.wallpaper.model.WallpaperColorsViewModel
@@ -50,6 +51,7 @@
     wallpaperPreviewNavigator: WallpaperPreviewNavigator,
     private val wallpaperInteractor: WallpaperInteractor,
     private val themedIconInteractor: ThemedIconInteractor,
+    private val colorPickerInteractor: ColorPickerInteractor,
     wallpaperManager: WallpaperManager,
     isTwoPaneAndSmallWidth: Boolean,
     savedInstanceState: Bundle?,
@@ -117,6 +119,7 @@
             initialExtrasProvider = { getInitialExtras(isOnLockScreen) },
             wallpaperInteractor = wallpaperInteractor,
             themedIconInteractor = themedIconInteractor,
+            colorPickerInteractor = colorPickerInteractor,
             screen = screen,
         )
     }
diff --git a/src/com/android/customization/picker/preview/ui/viewmodel/PreviewWithThemeViewModel.kt b/src/com/android/customization/picker/preview/ui/viewmodel/PreviewWithThemeViewModel.kt
index 435878d..83f986d 100644
--- a/src/com/android/customization/picker/preview/ui/viewmodel/PreviewWithThemeViewModel.kt
+++ b/src/com/android/customization/picker/preview/ui/viewmodel/PreviewWithThemeViewModel.kt
@@ -20,12 +20,14 @@
 import android.app.WallpaperColors
 import android.os.Bundle
 import com.android.customization.model.themedicon.domain.interactor.ThemedIconInteractor
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
 import com.android.wallpaper.model.WallpaperInfo
 import com.android.wallpaper.module.CustomizationSections
 import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
 import com.android.wallpaper.picker.customization.ui.viewmodel.ScreenPreviewViewModel
 import com.android.wallpaper.util.PreviewUtils
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.combine
 
 /** A ThemePicker version of the [ScreenPreviewViewModel] */
 class PreviewWithThemeViewModel(
@@ -35,6 +37,7 @@
     onWallpaperColorChanged: (WallpaperColors?) -> Unit = {},
     wallpaperInteractor: WallpaperInteractor,
     private val themedIconInteractor: ThemedIconInteractor? = null,
+    colorPickerInteractor: ColorPickerInteractor? = null,
     screen: CustomizationSections.Screen,
 ) :
     ScreenPreviewViewModel(
@@ -46,4 +49,16 @@
         screen,
     ) {
     override fun workspaceUpdateEvents(): Flow<Boolean>? = themedIconInteractor?.isActivated
+
+    private val wallpaperIsLoading = super.isLoading
+
+    override val isLoading: Flow<Boolean> =
+        colorPickerInteractor?.let {
+            combine(wallpaperIsLoading, colorPickerInteractor.isApplyingSystemColor) {
+                wallpaperIsLoading,
+                colorIsLoading ->
+                wallpaperIsLoading || colorIsLoading
+            }
+        }
+            ?: wallpaperIsLoading
 }