Merge "Fix discard changes when flex clock customization" into main
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 8fdc97e..f62a6ad 100755
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -204,6 +204,9 @@
     <!--Title for a grid option, describing the number of columns and rows, eg: 4x4 [CHAR LIMIT=10] -->
     <string name="grid_title_pattern"><xliff:g name="num_cols" example="1">%1$d</xliff:g>x<xliff:g name="num_rows" example="1">%2$d</xliff:g></string>
 
+    <!--Copntent description for a grid option, describing the number of columns and rows, eg: 4 by 4 [CHAR LIMIT=10] -->
+    <string name="grid_content_description"><xliff:g name="num_cols" example="1">%1$d</xliff:g> by <xliff:g name="num_rows" example="1">%2$d</xliff:g></string>
+
     <!-- Label of what will happen when user tap on apply button to change grid. [CHAR LIMIT=50] -->
     <string name="apply_grid_btn_note">Changing grid size will reload workspace and may take a few seconds.</string>
 
diff --git a/src/com/android/customization/model/color/ColorProvider.kt b/src/com/android/customization/model/color/ColorProvider.kt
index 74da5c2..ef5e3d6 100644
--- a/src/com/android/customization/model/color/ColorProvider.kt
+++ b/src/com/android/customization/model/color/ColorProvider.kt
@@ -120,15 +120,15 @@
             loaderJob?.join()
             if (presetColorBundles == null || reload) {
                 try {
-                    loaderJob = launch { loadPreset() }
+                    loaderJob = launch { loadPreset(isNewPickerUi) }
                 } catch (e: Throwable) {
                     colorsAvailable = false
                     callback?.onError(e)
                     return@launch
                 }
-                callback?.onOptionsLoaded(buildFinalList())
+                callback?.onOptionsLoaded(buildFinalList(isNewPickerUi))
             } else {
-                callback?.onOptionsLoaded(buildFinalList())
+                callback?.onOptionsLoaded(buildFinalList(isNewPickerUi))
             }
         }
     }
@@ -236,8 +236,8 @@
     }
 
     /**
-     * Returns the light theme version of the Revamped UI preview of a ColorScheme based on this
-     * order: top left, top right, bottom left, bottom right
+     * Returns the light theme preview of a dynamic ColorScheme based on this order: top left, top
+     * right, bottom left, bottom right
      *
      * This color mapping corresponds to GM3 colors: Primary (light), Primary (light), Secondary
      * LStar 85, and Tertiary LStar 70
@@ -253,8 +253,8 @@
     }
 
     /**
-     * Returns the dark theme version of the Revamped UI preview of a ColorScheme based on this
-     * order: top left, top right, bottom left, bottom right
+     * Returns the dark theme preview of a dynamic ColorScheme based on this order: top left, top
+     * right, bottom left, bottom right
      *
      * This color mapping corresponds to GM3 colors: Primary (dark), Primary (dark), Secondary LStar
      * 35, and Tertiary LStar 70
@@ -270,8 +270,8 @@
     }
 
     /**
-     * Returns the light theme version of the Revamped UI preview of a ColorScheme based on this
-     * order: top left, top right, bottom left, bottom right
+     * Returns the light theme preview of a monochrome ColorScheme based on this order: top left,
+     * top right, bottom left, bottom right
      *
      * This color mapping corresponds to GM3 colors: Primary LStar 0, Primary LStar 0, Secondary
      * LStar 85, and Tertiary LStar 70
@@ -287,8 +287,8 @@
     }
 
     /**
-     * Returns the dark theme version of the Revamped UI preview of a ColorScheme based on this
-     * order: top left, top right, bottom left, bottom right
+     * Returns the dark theme preview of a monochrome ColorScheme based on this order: top left, top
+     * right, bottom left, bottom right
      *
      * This color mapping corresponds to GM3 colors: Primary LStar 99, Primary LStar 99, Secondary
      * LStar 35, and Tertiary LStar 70
@@ -304,13 +304,13 @@
     }
 
     /**
-     * Returns the Revamped UI preview of a preset ColorScheme based on this order: top left, top
-     * right, bottom left, bottom right
+     * Returns the preview of a preset ColorScheme based on this order: top left, top right, bottom
+     * left, bottom right
      */
-    private fun getPresetColorPreview(colorScheme: ColorScheme, seed: Int): IntArray {
+    private fun getFixedPresetColorPreview(colorScheme: ColorScheme): IntArray {
         val colors =
             when (colorScheme.style) {
-                Style.FRUIT_SALAD -> intArrayOf(seed, colorScheme.accent1.s200)
+                Style.FRUIT_SALAD -> intArrayOf(colorScheme.accent3.s100, colorScheme.accent1.s200)
                 Style.TONAL_SPOT -> intArrayOf(colorScheme.accentColor, colorScheme.accentColor)
                 Style.RAINBOW -> intArrayOf(colorScheme.accent1.s200, colorScheme.accent1.s200)
                 else -> intArrayOf(colorScheme.accent1.s100, colorScheme.accent1.s100)
@@ -318,7 +318,30 @@
         return intArrayOf(colors[0], colors[1], colors[0], colors[1])
     }
 
-    private suspend fun loadPreset() =
+    /**
+     * Returns the light theme contrast-adjusted preview of a preset ColorScheme, specifically for
+     * Revamped UI, based on this order: top left, top right, bottom left, bottom right
+     */
+    private fun getLightPresetColorPreview(colorScheme: ColorScheme): IntArray {
+        val colors =
+            when (colorScheme.style) {
+                Style.FRUIT_SALAD ->
+                    intArrayOf(
+                        colorScheme.accent3.getAtTone(450f),
+                        colorScheme.accent1.getAtTone(550f),
+                    )
+                Style.TONAL_SPOT -> intArrayOf(colorScheme.accentColor, colorScheme.accentColor)
+                Style.RAINBOW ->
+                    intArrayOf(
+                        colorScheme.accent1.getAtTone(450f),
+                        colorScheme.accent1.getAtTone(450f),
+                    )
+                else -> intArrayOf(colorScheme.accent1.s100, colorScheme.accent1.s100)
+            }
+        return intArrayOf(colors[0], colors[1], colors[0], colors[1])
+    }
+
+    private suspend fun loadPreset(isNewPickerUi: Boolean) =
         withContext(Dispatchers.IO) {
             val bundles: MutableList<ColorOption> = ArrayList()
 
@@ -358,9 +381,23 @@
                         hasMonochrome = true
                         monochromeBundleName = bundleName
                     }
-                    bundles.add(buildPreset(bundleName, index, style))
+                    bundles.add(
+                        buildPreset(
+                            bundleName = bundleName,
+                            index = index,
+                            style = style,
+                            isNewPickerUi = isNewPickerUi,
+                        )
+                    )
                 } else {
-                    bundles.add(buildPreset(bundleName, index, null))
+                    bundles.add(
+                        buildPreset(
+                            bundleName = bundleName,
+                            index = index,
+                            style = null,
+                            isNewPickerUi = isNewPickerUi,
+                        )
+                    )
                 }
 
                 index++
@@ -378,6 +415,7 @@
         index: Int,
         @Style.Type style: Int? = null,
         type: ColorType = ColorType.PRESET_COLOR,
+        isNewPickerUi: Boolean,
     ): ColorOptionImpl {
         val builder = ColorOptionImpl.Builder()
         builder.title = getItemStringFromStub(COLOR_BUNDLE_NAME_PREFIX, bundleName)
@@ -406,8 +444,13 @@
                     lightColors = getLightMonochromePreview(lightColorScheme)
                 }
                 else -> {
-                    darkColors = getPresetColorPreview(darkColorScheme, colorFromStub)
-                    lightColors = getPresetColorPreview(lightColorScheme, colorFromStub)
+                    darkColors = getFixedPresetColorPreview(darkColorScheme)
+                    lightColors =
+                        if (isNewPickerUi) {
+                            getFixedPresetColorPreview(lightColorScheme)
+                        } else {
+                            getLightPresetColorPreview(lightColorScheme)
+                        }
                 }
             }
         }
@@ -416,7 +459,7 @@
         return builder.build()
     }
 
-    private fun buildFinalList(): List<ColorOption> {
+    private fun buildFinalList(isNewPickerUi: Boolean): List<ColorOption> {
         val presetColors = presetColorBundles ?: emptyList()
         val wallpaperColors = wallpaperColorBundles?.toMutableList() ?: mutableListOf()
         // Insert monochrome in the second position if it is enabled and included in preset
@@ -426,7 +469,13 @@
                 if (wallpaperColors.isNotEmpty()) {
                     wallpaperColors.add(
                         1,
-                        buildPreset(it, -1, Style.MONOCHROMATIC, ColorType.WALLPAPER_COLOR),
+                        buildPreset(
+                            bundleName = it,
+                            index = -1,
+                            style = Style.MONOCHROMATIC,
+                            type = ColorType.WALLPAPER_COLOR,
+                            isNewPickerUi = isNewPickerUi,
+                        ),
                     )
                 }
             }
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index b577e43..df68c32 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -70,6 +70,7 @@
 import com.android.wallpaper.module.PartnerProvider
 import com.android.wallpaper.module.WallpaperPicker2Injector
 import com.android.wallpaper.module.WallpaperPreferences
+import com.android.wallpaper.module.WallpaperRefresher
 import com.android.wallpaper.module.logging.UserEventLogger
 import com.android.wallpaper.network.Requester
 import com.android.wallpaper.picker.CustomizationPickerActivity
@@ -120,6 +121,7 @@
     wallpaperColorsRepository: Lazy<WallpaperColorsRepository>,
     defaultWallpaperCategoryWrapper: Lazy<WallpaperCategoryWrapper>,
     packageNotifier: Lazy<PackageStatusNotifier>,
+    wallpaperRefresher: Lazy<WallpaperRefresher>,
 ) :
     WallpaperPicker2Injector(
         mainScope,
@@ -135,6 +137,7 @@
         wallpaperColorsRepository,
         defaultWallpaperCategoryWrapper,
         packageNotifier,
+        wallpaperRefresher,
     ),
     CustomizationInjector {
     private var customizationSections: CustomizationSections? = null
diff --git a/src/com/android/customization/picker/grid/ui/viewmodel/GridScreenViewModel.kt b/src/com/android/customization/picker/grid/ui/viewmodel/GridScreenViewModel.kt
index 179127d..201b21f 100644
--- a/src/com/android/customization/picker/grid/ui/viewmodel/GridScreenViewModel.kt
+++ b/src/com/android/customization/picker/grid/ui/viewmodel/GridScreenViewModel.kt
@@ -26,6 +26,7 @@
 import com.android.customization.model.ResourceConstants
 import com.android.customization.picker.grid.domain.interactor.GridInteractor
 import com.android.customization.picker.grid.shared.model.GridOptionItemsModel
+import com.android.themepicker.R
 import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
 import kotlinx.coroutines.flow.Flow
@@ -34,10 +35,7 @@
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.launch
 
-class GridScreenViewModel(
-    context: Context,
-    private val interactor: GridInteractor,
-) : ViewModel() {
+class GridScreenViewModel(context: Context, private val interactor: GridInteractor) : ViewModel() {
 
     @SuppressLint("StaticFieldLeak") // We're not leaking this context as it is the app context.
     private val applicationContext = context.applicationContext
@@ -46,7 +44,7 @@
         interactor.options.map { model -> toViewModel(model) }
 
     private fun toViewModel(
-        model: GridOptionItemsModel,
+        model: GridOptionItemsModel
     ): List<OptionItemViewModel<GridIconViewModel>> {
         val iconShapePath =
             applicationContext.resources.getString(
@@ -72,6 +70,14 @@
                                 path = iconShapePath,
                             ),
                         text = text,
+                        contentDescription =
+                            Text.Loaded(
+                                applicationContext.resources.getString(
+                                    R.string.grid_content_description,
+                                    option.cols,
+                                    option.rows,
+                                )
+                            ),
                         isSelected = option.isSelected,
                         onClicked =
                             option.isSelected.map { isSelected ->
@@ -87,20 +93,14 @@
         }
     }
 
-    class Factory(
-        context: Context,
-        private val interactor: GridInteractor,
-    ) : ViewModelProvider.Factory {
+    class Factory(context: Context, private val interactor: GridInteractor) :
+        ViewModelProvider.Factory {
 
         private val applicationContext = context.applicationContext
 
         @Suppress("UNCHECKED_CAST")
         override fun <T : ViewModel> create(modelClass: Class<T>): T {
-            return GridScreenViewModel(
-                context = applicationContext,
-                interactor = interactor,
-            )
-                as T
+            return GridScreenViewModel(context = applicationContext, interactor = interactor) as T
         }
     }
 }
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt b/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
index 32ee209..704ab79 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
@@ -49,10 +49,11 @@
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.flow.stateIn
-import kotlinx.coroutines.flow.take
 
 class KeyguardQuickAffordancePickerViewModel2
 @AssistedInject
@@ -82,9 +83,15 @@
             )
     private val overridingQuickAffordances = MutableStateFlow<Map<String, String>>(emptyMap())
     private val selectedQuickAffordancesGroupBySlotId =
-        quickAffordanceInteractor.selections.map {
-            it.groupBy { selectionModel -> selectionModel.slotId }
-        }
+        quickAffordanceInteractor.selections
+            .map { it.groupBy { selectionModel -> selectionModel.slotId } }
+            .shareIn(viewModelScope, SharingStarted.WhileSubscribed(), 1)
+    // optimisticUpdateQuickAffordances updates right after applying button is clicked, while the
+    // actual update of selectedQuickAffordancesGroupBySlotId later updates until the system
+    // completes the update task. This can make sure the apply button state updates before we return
+    // to the previous screen.
+    private val optimisticUpdateQuickAffordances: MutableStateFlow<Map<String, String>?> =
+        MutableStateFlow(null)
 
     val previewingQuickAffordances =
         combine(
@@ -258,19 +265,18 @@
         }
 
     val onApply: Flow<(suspend () -> Unit)?> =
-        combine(overridingQuickAffordances, selectedQuickAffordancesGroupBySlotId) {
+        combine(overridingQuickAffordances, optimisticUpdateQuickAffordances) {
             overridingQuickAffordances,
-            selectedQuickAffordancesGroupBySlotId ->
-            // If all overridingQuickAffordances are same as the selected quick affordances, it is
-            // not yet edited
+            optimisticUpdateQuickAffordances ->
+            // If all overridingQuickAffordances is empty or are same as the
+            // optimisticUpdateQuickAffordances, it is not yet edited
             val isQuickAffordancesEdited =
                 (!overridingQuickAffordances.all { (slotId, overridingQuickAffordanceId) ->
-                    selectedQuickAffordancesGroupBySlotId[slotId]?.find {
-                        it.affordanceId == overridingQuickAffordanceId
-                    } != null
+                    optimisticUpdateQuickAffordances?.get(slotId) == overridingQuickAffordanceId
                 })
             if (isQuickAffordancesEdited) {
                 {
+                    this.optimisticUpdateQuickAffordances.value = overridingQuickAffordances
                     overridingQuickAffordances.forEach { entry ->
                         val slotId = entry.key
                         val affordanceId = entry.value
@@ -282,10 +288,6 @@
                                 affordanceId = affordanceId,
                             )
                         }
-                        // Suspend until the next selectedQuickAffordancesGroupBySlotId update
-                        this.selectedQuickAffordancesGroupBySlotId.take(1).collect {
-                            return@collect
-                        }
                         logger.logShortcutApplied(shortcut = affordanceId, shortcutSlotId = slotId)
                     }
                 }