Merge "Fix missing color option on activity resume (1/2)" into udc-dev am: 79ffe8d5bf

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

Change-Id: Ib57061f4b6df9ffe30153727ef24512d7761f972
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

3 files changed, 142 insertions, 58 deletions

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 464e8bf9..41ef3a57 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
@@ -16,7 +16,6 @@
  */
 package com.android.customization.picker.color.data.repository
 
-import android.app.WallpaperColors
 import android.util.Log
 import com.android.customization.model.CustomizationManager
 import com.android.customization.model.color.ColorCustomizationManager
@@ -25,6 +24,7 @@ import com.android.customization.model.color.ColorOptionImpl
 import com.android.customization.picker.color.shared.model.ColorOptionModel
 import com.android.customization.picker.color.shared.model.ColorType
 import com.android.systemui.monet.Style
+import com.android.wallpaper.model.WallpaperColorsModel
 import com.android.wallpaper.model.WallpaperColorsViewModel
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
@@ -39,9 +39,9 @@ class ColorPickerRepositoryImpl(
     private val colorManager: ColorCustomizationManager,
 ) : ColorPickerRepository {
 
-    private val homeWallpaperColors: StateFlow<WallpaperColors?> =
+    private val homeWallpaperColors: StateFlow<WallpaperColorsModel?> =
         wallpaperColorsViewModel.homeWallpaperColors
-    private val lockWallpaperColors: StateFlow<WallpaperColors?> =
+    private val lockWallpaperColors: StateFlow<WallpaperColorsModel?> =
         wallpaperColorsViewModel.lockWallpaperColors
 
     override val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>> =
@@ -50,7 +50,26 @@ class ColorPickerRepositoryImpl(
             }
             .map { (homeColors, lockColors) ->
                 suspendCancellableCoroutine { continuation ->
-                    colorManager.setWallpaperColors(homeColors, lockColors)
+                    if (
+                        homeColors is WallpaperColorsModel.Loading ||
+                            lockColors is WallpaperColorsModel.Loading
+                    ) {
+                        continuation.resumeWith(
+                            Result.success(
+                                mapOf(
+                                    ColorType.WALLPAPER_COLOR to listOf(),
+                                    ColorType.PRESET_COLOR to listOf()
+                                )
+                            )
+                        )
+                        return@suspendCancellableCoroutine
+                    }
+                    val homeColorsLoaded = homeColors as WallpaperColorsModel.Loaded
+                    val lockColorsLoaded = lockColors as WallpaperColorsModel.Loaded
+                    colorManager.setWallpaperColors(
+                        homeColorsLoaded.colors,
+                        lockColorsLoaded.colors
+                    )
                     colorManager.fetchRevampedUIOptions(
                         object : CustomizationManager.OptionsFetchedListener<ColorOption?> {
                             override fun onOptionsLoaded(options: MutableList<ColorOption?>?) {
diff --git a/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt b/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
index e322eac4..78bfa43e 100644
--- a/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
+++ b/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
@@ -15,6 +15,7 @@
  */
 package com.android.customization.picker.color.ui.fragment
 
+import android.app.WallpaperManager
 import android.os.Bundle
 import android.view.LayoutInflater
 import android.view.View
@@ -23,10 +24,13 @@ import android.widget.FrameLayout
 import androidx.cardview.widget.CardView
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.get
+import androidx.lifecycle.lifecycleScope
 import com.android.customization.model.mode.DarkModeSectionController
 import com.android.customization.module.ThemePickerInjector
 import com.android.customization.picker.color.ui.binder.ColorPickerBinder
 import com.android.wallpaper.R
+import com.android.wallpaper.model.WallpaperColorsModel
+import com.android.wallpaper.model.WallpaperColorsViewModel
 import com.android.wallpaper.module.CustomizationSections
 import com.android.wallpaper.module.InjectorProvider
 import com.android.wallpaper.picker.AppbarFragment
@@ -34,8 +38,11 @@ import com.android.wallpaper.picker.customization.ui.binder.ScreenPreviewBinder
 import com.android.wallpaper.picker.customization.ui.viewmodel.ScreenPreviewViewModel
 import com.android.wallpaper.util.DisplayUtils
 import com.android.wallpaper.util.PreviewUtils
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.suspendCancellableCoroutine
+import kotlinx.coroutines.withContext
 
 @OptIn(ExperimentalCoroutinesApi::class)
 class ColorPickerFragment : AppbarFragment() {
@@ -66,6 +73,7 @@ class ColorPickerFragment : AppbarFragment() {
         val wallpaperInfoFactory = injector.getCurrentWallpaperInfoFactory(requireContext())
         val displayUtils: DisplayUtils = injector.getDisplayUtils(requireContext())
         val wcViewModel = injector.getWallpaperColorsViewModel()
+        val wallpaperManager = WallpaperManager.getInstance(requireContext())
 
         binding =
             ColorPickerBinder.bind(
@@ -102,6 +110,18 @@ class ColorPickerFragment : AppbarFragment() {
                         suspendCancellableCoroutine { continuation ->
                             wallpaperInfoFactory.createCurrentWallpaperInfos(
                                 { homeWallpaper, lockWallpaper, _ ->
+                                    lifecycleScope.launch {
+                                        if (
+                                            wcViewModel.lockWallpaperColors.value
+                                                is WallpaperColorsModel.Loading
+                                        ) {
+                                            loadInitialColors(
+                                                wallpaperManager,
+                                                wcViewModel,
+                                                CustomizationSections.Screen.LOCK_SCREEN
+                                            )
+                                        }
+                                    }
                                     continuation.resume(lockWallpaper ?: homeWallpaper, null)
                                 },
                                 forceReload,
@@ -137,6 +157,18 @@ class ColorPickerFragment : AppbarFragment() {
                         suspendCancellableCoroutine { continuation ->
                             wallpaperInfoFactory.createCurrentWallpaperInfos(
                                 { homeWallpaper, lockWallpaper, _ ->
+                                    lifecycleScope.launch {
+                                        if (
+                                            wcViewModel.homeWallpaperColors.value
+                                                is WallpaperColorsModel.Loading
+                                        ) {
+                                            loadInitialColors(
+                                                wallpaperManager,
+                                                wcViewModel,
+                                                CustomizationSections.Screen.HOME_SCREEN
+                                            )
+                                        }
+                                    }
                                     continuation.resume(homeWallpaper ?: lockWallpaper, null)
                                 },
                                 forceReload,
@@ -144,7 +176,7 @@ class ColorPickerFragment : AppbarFragment() {
                         }
                     },
                     onWallpaperColorChanged = { colors ->
-                        wcViewModel.setLockWallpaperColors(colors)
+                        wcViewModel.setHomeWallpaperColors(colors)
                     },
                     wallpaperInteractor = injector.getWallpaperInteractor(requireContext()),
                     screen = CustomizationSections.Screen.HOME_SCREEN,
@@ -168,6 +200,30 @@ class ColorPickerFragment : AppbarFragment() {
         return view
     }
 
+    private suspend fun loadInitialColors(
+        wallpaperManager: WallpaperManager,
+        colorViewModel: WallpaperColorsViewModel,
+        screen: CustomizationSections.Screen,
+    ) {
+        withContext(Dispatchers.IO) {
+            val colors =
+                wallpaperManager.getWallpaperColors(
+                    if (screen == CustomizationSections.Screen.LOCK_SCREEN) {
+                        WallpaperManager.FLAG_LOCK
+                    } else {
+                        WallpaperManager.FLAG_SYSTEM
+                    }
+                )
+            withContext(Dispatchers.Main) {
+                if (screen == CustomizationSections.Screen.LOCK_SCREEN) {
+                    colorViewModel.setLockWallpaperColors(colors)
+                } else {
+                    colorViewModel.setHomeWallpaperColors(colors)
+                }
+            }
+        }
+    }
+
     override fun onSaveInstanceState(savedInstanceState: Bundle) {
         super.onSaveInstanceState(savedInstanceState)
         binding?.saveInstanceState(savedInstanceState)
diff --git a/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt b/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
index 8bb128fd..67c68387 100644
--- a/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
+++ b/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
@@ -93,60 +93,69 @@ private constructor(
     /** The list of all color options mapped by their color type */
     private val allColorOptions:
         Flow<Map<ColorType, List<OptionItemViewModel<ColorOptionIconViewModel>>>> =
-        interactor.colorOptions.map { colorOptions ->
-            colorOptions
-                .map { colorOptionEntry ->
-                    colorOptionEntry.key to
-                        colorOptionEntry.value.map { colorOptionModel ->
-                            val colorOption: ColorOptionImpl =
-                                colorOptionModel.colorOption as ColorOptionImpl
-                            val lightThemeColors =
-                                colorOption.previewInfo.resolveColors(/* darkTheme= */ false)
-                            val darkThemeColors =
-                                colorOption.previewInfo.resolveColors(/* darkTheme= */ true)
-                            val isSelectedFlow: StateFlow<Boolean> =
-                                interactor.activeColorOption
-                                    .map {
-                                        it?.colorOption?.isEquivalent(colorOptionModel.colorOption)
-                                            ?: colorOptionModel.isSelected
-                                    }
-                                    .stateIn(viewModelScope)
-                            OptionItemViewModel<ColorOptionIconViewModel>(
-                                key = MutableStateFlow(colorOptionModel.key) as StateFlow<String>,
-                                payload =
-                                    ColorOptionIconViewModel(
-                                        lightThemeColor0 = lightThemeColors[0],
-                                        lightThemeColor1 = lightThemeColors[1],
-                                        lightThemeColor2 = lightThemeColors[2],
-                                        lightThemeColor3 = lightThemeColors[3],
-                                        darkThemeColor0 = darkThemeColors[0],
-                                        darkThemeColor1 = darkThemeColors[1],
-                                        darkThemeColor2 = darkThemeColors[2],
-                                        darkThemeColor3 = darkThemeColors[3],
-                                    ),
-                                text =
-                                    Text.Loaded(
-                                        colorOption.getContentDescription(context).toString()
-                                    ),
-                                isTextUserVisible = false,
-                                isSelected = isSelectedFlow,
-                                onClicked =
-                                    isSelectedFlow.map { isSelected ->
-                                        if (isSelected) {
-                                            null
-                                        } else {
-                                            {
-                                                viewModelScope.launch {
-                                                    interactor.select(colorOptionModel)
+        interactor.colorOptions
+            .map { colorOptions ->
+                colorOptions
+                    .map { colorOptionEntry ->
+                        colorOptionEntry.key to
+                            colorOptionEntry.value.map { colorOptionModel ->
+                                val colorOption: ColorOptionImpl =
+                                    colorOptionModel.colorOption as ColorOptionImpl
+                                val lightThemeColors =
+                                    colorOption.previewInfo.resolveColors(/* darkTheme= */ false)
+                                val darkThemeColors =
+                                    colorOption.previewInfo.resolveColors(/* darkTheme= */ true)
+                                val isSelectedFlow: StateFlow<Boolean> =
+                                    interactor.activeColorOption
+                                        .map {
+                                            it?.colorOption?.isEquivalent(
+                                                colorOptionModel.colorOption
+                                            )
+                                                ?: colorOptionModel.isSelected
+                                        }
+                                        .stateIn(viewModelScope)
+                                OptionItemViewModel<ColorOptionIconViewModel>(
+                                    key =
+                                        MutableStateFlow(colorOptionModel.key) as StateFlow<String>,
+                                    payload =
+                                        ColorOptionIconViewModel(
+                                            lightThemeColor0 = lightThemeColors[0],
+                                            lightThemeColor1 = lightThemeColors[1],
+                                            lightThemeColor2 = lightThemeColors[2],
+                                            lightThemeColor3 = lightThemeColors[3],
+                                            darkThemeColor0 = darkThemeColors[0],
+                                            darkThemeColor1 = darkThemeColors[1],
+                                            darkThemeColor2 = darkThemeColors[2],
+                                            darkThemeColor3 = darkThemeColors[3],
+                                        ),
+                                    text =
+                                        Text.Loaded(
+                                            colorOption.getContentDescription(context).toString()
+                                        ),
+                                    isTextUserVisible = false,
+                                    isSelected = isSelectedFlow,
+                                    onClicked =
+                                        isSelectedFlow.map { isSelected ->
+                                            if (isSelected) {
+                                                null
+                                            } else {
+                                                {
+                                                    viewModelScope.launch {
+                                                        interactor.select(colorOptionModel)
+                                                    }
                                                 }
                                             }
-                                        }
-                                    },
-                            )
-                        }
-                }
-                .toMap()
-        }
+                                        },
+                                )
+                            }
+                    }
+                    .toMap()
+            }
+            .shareIn(
+                scope = viewModelScope,
+                started = SharingStarted.WhileSubscribed(),
+                replay = 1,
+            )
 
     /** The list of all available color options for the selected Color Type. */
     val colorOptions: Flow<List<OptionItemViewModel<ColorOptionIconViewModel>>> =
@@ -159,7 +168,7 @@ private constructor(
             }
             .shareIn(
                 scope = viewModelScope,
-                started = SharingStarted.WhileSubscribed(),
+                started = SharingStarted.Eagerly,
                 replay = 1,
             )