Merge "Clean up and refactor load initial colors (1/3)" into udc-dev

10 files changed, 62 insertions, 40 deletions

diff --git a/src/com/android/customization/module/DefaultCustomizationSections.java b/src/com/android/customization/module/DefaultCustomizationSections.java
index b408a89d..00f9433c 100644
--- a/src/com/android/customization/module/DefaultCustomizationSections.java
+++ b/src/com/android/customization/module/DefaultCustomizationSections.java
@@ -145,6 +145,7 @@ public final class DefaultCustomizationSections implements CustomizationSections
                                         activity,
                                         mColorPickerViewModelFactory)
                                         .get(ColorPickerViewModel.class),
+                                wallpaperColorsViewModel,
                                 lifecycleOwner),
                         // Wallpaper quick switch section.
                         new WallpaperQuickSwitchSectionController(
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index d00ed28e..1e0dbe5f 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -439,7 +439,8 @@ open class ThemePickerInjector : WallpaperPicker2Injector(), CustomizationInject
                     repository =
                         ColorPickerRepositoryImpl(
                             wallpaperColorsViewModel,
-                            getColorCustomizationManager(appContext)
+                            getColorCustomizationManager(appContext),
+                            WallpaperManager.getInstance(appContext),
                         ),
                     snapshotRestorer = {
                         getColorPickerSnapshotRestorer(appContext, wallpaperColorsViewModel)
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 7cf9fd03..ba77f25f 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
@@ -37,4 +37,7 @@ interface ColorPickerRepository {
 
     /** Returns the current selected color source based on system settings */
     fun getCurrentColorSource(): String?
+
+    /** Retrieves and stores the wallpaper colors for generating wallpaper color options */
+    suspend fun loadInitialColors()
 }
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 41ef3a57..f1c695dd 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
@@ -16,6 +16,7 @@
  */
 package com.android.customization.picker.color.data.repository
 
+import android.app.WallpaperManager
 import android.util.Log
 import com.android.customization.model.CustomizationManager
 import com.android.customization.model.color.ColorCustomizationManager
@@ -26,17 +27,20 @@ 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.Dispatchers
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.suspendCancellableCoroutine
+import kotlinx.coroutines.withContext
 
 // TODO (b/262924623): refactor to remove dependency on ColorCustomizationManager & ColorOption
 // TODO (b/268203200): Create test for ColorPickerRepositoryImpl
 class ColorPickerRepositoryImpl(
-    wallpaperColorsViewModel: WallpaperColorsViewModel,
+    private val wallpaperColorsViewModel: WallpaperColorsViewModel,
     private val colorManager: ColorCustomizationManager,
+    private val wallpaperManager: WallpaperManager,
 ) : ColorPickerRepository {
 
     private val homeWallpaperColors: StateFlow<WallpaperColorsModel?> =
@@ -151,6 +155,17 @@ class ColorPickerRepositoryImpl(
         return colorManager.currentColorSource
     }
 
+    override suspend fun loadInitialColors() {
+        withContext(Dispatchers.IO) {
+            val lockColors = wallpaperManager.getWallpaperColors(WallpaperManager.FLAG_LOCK)
+            val homeColors = wallpaperManager.getWallpaperColors(WallpaperManager.FLAG_SYSTEM)
+            withContext(Dispatchers.Main) {
+                wallpaperColorsViewModel.setLockWallpaperColors(lockColors)
+                wallpaperColorsViewModel.setHomeWallpaperColors(homeColors)
+            }
+        }
+    }
+
     private fun ColorOptionImpl.toModel(): ColorOptionModel {
         return ColorOptionModel(
             key = "${this.type}::${this.style}::${this.serializedPackages}",
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 714129df..e4b8795e 100644
--- a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
@@ -166,6 +166,8 @@ class FakeColorPickerRepository(private val context: Context) : ColorPickerRepos
             else -> null
         }
 
+    override suspend fun loadInitialColors() {}
+
     private fun ColorOptionModel.testEquals(other: Any?): Boolean {
         if (other == null) {
             return false
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 8c7a4b72..36d2a0d4 100644
--- a/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt
+++ b/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt
@@ -46,4 +46,6 @@ class ColorPickerInteractor(
     }
 
     fun getCurrentColorOption(): ColorOptionModel = repository.getCurrentColorOption()
+
+    suspend fun loadInitialColors() = repository.loadInitialColors()
 }
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 78bfa43e..3390a0d1 100644
--- a/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
+++ b/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
@@ -28,6 +28,7 @@ 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.customization.picker.color.ui.viewmodel.ColorPickerViewModel
 import com.android.wallpaper.R
 import com.android.wallpaper.model.WallpaperColorsModel
 import com.android.wallpaper.model.WallpaperColorsViewModel
@@ -73,20 +74,29 @@ class ColorPickerFragment : AppbarFragment() {
         val wallpaperInfoFactory = injector.getCurrentWallpaperInfoFactory(requireContext())
         val displayUtils: DisplayUtils = injector.getDisplayUtils(requireContext())
         val wcViewModel = injector.getWallpaperColorsViewModel()
-        val wallpaperManager = WallpaperManager.getInstance(requireContext())
+
+        val colorPickerViewModel =
+            ViewModelProvider(
+                    requireActivity(),
+                    injector.getColorPickerViewModelFactory(
+                        context = requireContext(),
+                        wallpaperColorsViewModel = wcViewModel,
+                    ),
+                )
+                .get() as ColorPickerViewModel
+
+        // load wallpaper colors if it has not been populated
+        if (
+            wcViewModel.lockWallpaperColors.value is WallpaperColorsModel.Loading ||
+                wcViewModel.homeWallpaperColors.value is WallpaperColorsModel.Loading
+        ) {
+            lifecycleScope.launch { colorPickerViewModel.loadInitialColors() }
+        }
 
         binding =
             ColorPickerBinder.bind(
                 view = view,
-                viewModel =
-                    ViewModelProvider(
-                            requireActivity(),
-                            injector.getColorPickerViewModelFactory(
-                                context = requireContext(),
-                                wallpaperColorsViewModel = wcViewModel,
-                            ),
-                        )
-                        .get(),
+                viewModel = colorPickerViewModel,
                 lifecycleOwner = this,
             )
 
@@ -110,18 +120,6 @@ 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,
@@ -157,18 +155,6 @@ 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,
diff --git a/src/com/android/customization/picker/color/ui/section/ColorSectionController2.kt b/src/com/android/customization/picker/color/ui/section/ColorSectionController2.kt
index f1c982b4..d0ddb2b0 100644
--- a/src/com/android/customization/picker/color/ui/section/ColorSectionController2.kt
+++ b/src/com/android/customization/picker/color/ui/section/ColorSectionController2.kt
@@ -20,6 +20,7 @@ package com.android.customization.picker.color.ui.section
 import android.content.Context
 import android.view.LayoutInflater
 import androidx.lifecycle.LifecycleOwner
+import androidx.lifecycle.lifecycleScope
 import com.android.customization.picker.color.ui.binder.ColorSectionViewBinder
 import com.android.customization.picker.color.ui.fragment.ColorPickerFragment
 import com.android.customization.picker.color.ui.view.ColorSectionView2
@@ -27,10 +28,14 @@ import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
 import com.android.wallpaper.R
 import com.android.wallpaper.model.CustomizationSectionController
 import com.android.wallpaper.model.CustomizationSectionController.CustomizationSectionNavigationController as NavigationController
+import com.android.wallpaper.model.WallpaperColorsModel
+import com.android.wallpaper.model.WallpaperColorsViewModel
+import kotlinx.coroutines.launch
 
 class ColorSectionController2(
     private val navigationController: NavigationController,
     private val viewModel: ColorPickerViewModel,
+    private val wcViewModel: WallpaperColorsViewModel,
     private val lifecycleOwner: LifecycleOwner
 ) : CustomizationSectionController<ColorSectionView2> {
 
@@ -53,6 +58,15 @@ class ColorSectionController2(
                     R.layout.color_section_view2,
                     null,
                 ) as ColorSectionView2
+
+        // load wallpaper colors if it has not been populated
+        if (
+            wcViewModel.lockWallpaperColors.value is WallpaperColorsModel.Loading ||
+                wcViewModel.homeWallpaperColors.value is WallpaperColorsModel.Loading
+        ) {
+            lifecycleOwner.lifecycleScope.launch { viewModel.loadInitialColors() }
+        }
+
         ColorSectionViewBinder.bind(
             view = view,
             viewModel = viewModel,
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 67c68387..057eff5d 100644
--- a/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
+++ b/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
@@ -202,6 +202,8 @@ private constructor(
                 replay = 1,
             )
 
+    suspend fun loadInitialColors() = interactor.loadInitialColors()
+
     class Factory(
         private val context: Context,
         private val interactor: ColorPickerInteractor,
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 56c6c30a..b6301225 100644
--- a/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
+++ b/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
@@ -94,10 +94,6 @@ open class PreviewWithThemeSectionController(
                                 } else {
                                     homeWallpaper ?: lockWallpaper
                                 }
-                            loadInitialColors(
-                                context = context,
-                                screen = screen,
-                            )
                             continuation.resume(wallpaper, null)
                         },
                         forceReload,