Merge "Import translations. DO NOT MERGE ANYWHERE" into udc-dev
diff --git a/src/com/android/customization/model/color/ColorOptionImpl.kt b/src/com/android/customization/model/color/ColorOptionImpl.kt
new file mode 100644
index 0000000..70aaa92
--- /dev/null
+++ b/src/com/android/customization/model/color/ColorOptionImpl.kt
@@ -0,0 +1,110 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+package com.android.customization.model.color
+
+import android.content.Context
+import android.view.View
+import androidx.annotation.ColorInt
+import com.android.customization.model.color.ColorOptionsProvider.ColorSource
+import com.android.customization.picker.color.shared.model.ColorType
+import com.android.systemui.monet.Style
+import com.android.wallpaper.R
+
+/**
+ * Represents a color option in the revamped UI, it can be used for both wallpaper and preset colors
+ */
+class ColorOptionImpl(
+    title: String?,
+    overlayPackages: Map<String, String?>,
+    isDefault: Boolean,
+    private val source: String?,
+    style: Style,
+    index: Int,
+    private val previewInfo: PreviewInfo,
+    val type: ColorType,
+) : ColorOption(title, overlayPackages, isDefault, style, index) {
+
+    class PreviewInfo(
+        @ColorInt val lightColors: IntArray,
+        @ColorInt val darkColors: IntArray,
+    ) : ColorOption.PreviewInfo {
+        @ColorInt
+        fun resolveColors(darkTheme: Boolean): IntArray {
+            return if (darkTheme) darkColors else lightColors
+        }
+    }
+
+    override fun bindThumbnailTile(view: View?) {
+        // Do nothing. This function will no longer be used in the Revamped UI
+    }
+
+    override fun getLayoutResId(): Int {
+        return R.layout.color_option
+    }
+
+    override fun getPreviewInfo(): PreviewInfo {
+        return previewInfo
+    }
+
+    override fun getContentDescription(context: Context): CharSequence? {
+        if (type == ColorType.WALLPAPER_COLOR) {
+            return context.getString(R.string.wallpaper_color_title)
+        }
+        return super.getContentDescription(context)
+    }
+
+    override fun getSource(): String? {
+        return source
+    }
+
+    class Builder {
+        var title: String? = null
+
+        @ColorInt var lightColors: IntArray = intArrayOf()
+
+        @ColorInt var darkColors: IntArray = intArrayOf()
+
+        @ColorSource var source: String? = null
+        var isDefault = false
+        var style = Style.TONAL_SPOT
+        var index = 0
+        var packages: MutableMap<String, String?> = HashMap()
+        var type = ColorType.WALLPAPER_COLOR
+
+        fun build(): ColorOptionImpl {
+            return ColorOptionImpl(
+                title,
+                packages,
+                isDefault,
+                source,
+                style,
+                index,
+                createPreviewInfo(),
+                type
+            )
+        }
+
+        private fun createPreviewInfo(): PreviewInfo {
+            return PreviewInfo(lightColors, darkColors)
+        }
+
+        fun addOverlayPackage(category: String?, packageName: String?): ColorOptionImpl.Builder {
+            category?.let { packages[category] = packageName }
+            return this
+        }
+    }
+}
diff --git a/src/com/android/customization/model/color/ColorProvider.kt b/src/com/android/customization/model/color/ColorProvider.kt
index 639a886..6520b52 100644
--- a/src/com/android/customization/model/color/ColorProvider.kt
+++ b/src/com/android/customization/model/color/ColorProvider.kt
@@ -34,11 +34,11 @@
 import com.android.customization.model.color.ColorOptionsProvider.COLOR_SOURCE_HOME
 import com.android.customization.model.color.ColorOptionsProvider.COLOR_SOURCE_LOCK
 import com.android.customization.model.color.ColorUtils.toColorString
+import com.android.customization.picker.color.shared.model.ColorType
 import com.android.systemui.monet.ColorScheme
 import com.android.systemui.monet.Style
 import com.android.wallpaper.compat.WallpaperManagerCompat
 import com.android.wallpaper.module.InjectorProvider
-import java.util.*
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.SupervisorJob
@@ -67,6 +67,10 @@
             arrayOf(Style.TONAL_SPOT, Style.SPRITZ, Style.VIBRANT, Style.EXPRESSIVE)
         else arrayOf(Style.TONAL_SPOT)
 
+    private val monochromeEnabled =
+        InjectorProvider.getInjector().getFlags().isMonochromaticThemeEnabled(mContext)
+    private var monochromeBundleName: String? = null
+
     private val scope =
         if (mContext is LifecycleOwner) {
             mContext.lifecycleScope
@@ -101,7 +105,7 @@
             scope.launch {
                 try {
                     if (colorBundles == null || reload) {
-                        loadPreset()
+                        loadPreset(shouldUseRevampedUi)
                     }
                     if (wallpaperColorsChanged || reload) {
                         loadSeedColors(
@@ -176,7 +180,31 @@
             )
         }
 
-        bundles.addAll(colorBundles?.filterNot { it is ColorSeedOption } ?: emptyList())
+        if (shouldUseRevampedUi) {
+            // Insert monochrome in the second position if it is enabled and included in preset
+            // colors
+            if (monochromeEnabled) {
+                monochromeBundleName?.let {
+                    bundles.add(
+                        1,
+                        buildRevampedUIPreset(
+                            it,
+                            -1,
+                            Style.MONOCHROMATIC,
+                            ColorType.WALLPAPER_COLOR
+                        )
+                    )
+                }
+            }
+            bundles.addAll(
+                colorBundles?.filterNot {
+                    (it as ColorOptionImpl).type == ColorType.WALLPAPER_COLOR
+                }
+                    ?: emptyList()
+            )
+        } else {
+            bundles.addAll(colorBundles?.filterNot { it is ColorSeedOption } ?: emptyList())
+        }
         colorBundles = bundles
     }
 
@@ -205,41 +233,56 @@
         shouldUseRevampedUi: Boolean,
     ) {
         // TODO(b/202145216): Measure time cost in the loop.
-        for (style in styleList) {
-            val builder = ColorSeedOption.Builder()
-            val lightColorScheme = ColorScheme(colorInt, /* darkTheme= */ false, style)
-            val darkColorScheme = ColorScheme(colorInt, /* darkTheme= */ true, style)
-            builder
-                .setLightColors(
-                    if (shouldUseRevampedUi) lightColorScheme.getRevampedUILightColorPreview()
-                    else lightColorScheme.getLightColorPreview()
-                )
-                .setDarkColors(
-                    if (shouldUseRevampedUi) darkColorScheme.getRevampedUIDarkColorPreview()
-                    else darkColorScheme.getDarkColorPreview()
-                )
-                .addOverlayPackage(
+        if (shouldUseRevampedUi) {
+            for (style in styleList) {
+                val lightColorScheme = ColorScheme(colorInt, /* darkTheme= */ false, style)
+                val darkColorScheme = ColorScheme(colorInt, /* darkTheme= */ true, style)
+                val builder = ColorOptionImpl.Builder()
+                builder.lightColors = getRevampedUILightColorPreview(lightColorScheme)
+                builder.darkColors = getRevampedUIDarkColorPreview(darkColorScheme)
+                builder.addOverlayPackage(
                     OVERLAY_CATEGORY_SYSTEM_PALETTE,
                     if (isDefault) "" else toColorString(colorInt)
                 )
-                .addOverlayPackage(
-                    OVERLAY_CATEGORY_COLOR,
-                    if (isDefault) "" else toColorString(colorInt)
-                )
-                .setSource(source)
-                .setStyle(style)
+                builder.source = source
+                builder.style = style
                 // Color option index value starts from 1.
-                .setIndex(i + 1)
+                builder.index = i + 1
+                builder.isDefault = isDefault
+                builder.type = ColorType.WALLPAPER_COLOR
+                bundles.add(builder.build())
+            }
+        } else {
+            for (style in styleList) {
+                val lightColorScheme = ColorScheme(colorInt, /* darkTheme= */ false, style)
+                val darkColorScheme = ColorScheme(colorInt, /* darkTheme= */ true, style)
+                val builder = ColorSeedOption.Builder()
+                builder
+                    .setLightColors(lightColorScheme.getLightColorPreview())
+                    .setDarkColors(darkColorScheme.getDarkColorPreview())
+                    .addOverlayPackage(
+                        OVERLAY_CATEGORY_SYSTEM_PALETTE,
+                        if (isDefault) "" else toColorString(colorInt)
+                    )
+                    .addOverlayPackage(
+                        OVERLAY_CATEGORY_COLOR,
+                        if (isDefault) "" else toColorString(colorInt)
+                    )
+                    .setSource(source)
+                    .setStyle(style)
+                    // Color option index value starts from 1.
+                    .setIndex(i + 1)
 
-            if (isDefault) builder.asDefault()
+                if (isDefault) builder.asDefault()
 
-            bundles.add(builder.build())
+                bundles.add(builder.build())
+            }
         }
     }
 
     /**
      * Returns the colors for the light theme version of the preview of a ColorScheme based on this
-     * order: |-------| | 0 | 1 | |---+---| | 2 | 3 | |-------|
+     * order: top left, top right, bottom left, bottom right
      */
     @ColorInt
     private fun ColorScheme.getLightColorPreview(): IntArray {
@@ -263,34 +306,57 @@
 
     /**
      * Returns the color for the dark theme version of the preview of a ColorScheme based on this
-     * order: |-------| | 0 | 1 | |---+---| | 2 | 3 | |-------|
+     * order: top left, top right, bottom left, bottom right
      */
     @ColorInt
     private fun ColorScheme.getDarkColorPreview(): IntArray {
         return getLightColorPreview()
     }
 
+    /**
+     * 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
+     */
     @ColorInt
-    private fun ColorScheme.getRevampedUILightColorPreview(): IntArray {
+    private fun getRevampedUILightColorPreview(colorScheme: ColorScheme): IntArray {
         return intArrayOf(
-            setAlphaComponent(this.accent1.s600, ALPHA_MASK),
-            setAlphaComponent(this.accent1.s600, ALPHA_MASK),
-            setAlphaComponent(this.accent2.s100, ALPHA_MASK),
-            ColorStateList.valueOf(this.accent3.s500).withLStar(59f).colors[0],
+            setAlphaComponent(colorScheme.accent1.s600, ALPHA_MASK),
+            setAlphaComponent(colorScheme.accent1.s600, ALPHA_MASK),
+            setAlphaComponent(colorScheme.accent2.s100, ALPHA_MASK),
+            ColorStateList.valueOf(colorScheme.accent3.s500).withLStar(59f).colors[0],
         )
     }
 
     /**
-     * Returns the color for the dark theme version of the preview of a ColorScheme based on this
-     * order: |-------| | 0 | 1 | |---+---| | 2 | 3 | |-------|
+     * 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
      */
     @ColorInt
-    private fun ColorScheme.getRevampedUIDarkColorPreview(): IntArray {
+    private fun getRevampedUIDarkColorPreview(colorScheme: ColorScheme): IntArray {
         return intArrayOf(
-            setAlphaComponent(this.accent1.s200, ALPHA_MASK),
-            setAlphaComponent(this.accent1.s200, ALPHA_MASK),
-            setAlphaComponent(this.accent2.s700, ALPHA_MASK),
-            setAlphaComponent(this.accent3.s100, ALPHA_MASK),
+            setAlphaComponent(colorScheme.accent1.s200, ALPHA_MASK),
+            setAlphaComponent(colorScheme.accent1.s200, ALPHA_MASK),
+            setAlphaComponent(colorScheme.accent2.s700, ALPHA_MASK),
+            setAlphaComponent(colorScheme.accent3.s100, ALPHA_MASK),
+        )
+    }
+
+    /**
+     * Returns the Revamped UI preview of a preset ColorScheme based on this order: top left, top
+     * right, bottom left, bottom right
+     */
+    private fun getRevampedUIPresetColorPreview(colorScheme: ColorScheme, seed: Int): IntArray {
+        val colors =
+            when (colorScheme.style) {
+                Style.FRUIT_SALAD -> intArrayOf(seed, colorScheme.accent1.s100)
+                Style.TONAL_SPOT -> intArrayOf(colorScheme.accentColor, colorScheme.accentColor)
+                else -> intArrayOf(colorScheme.accent1.s100, colorScheme.accent1.s100)
+            }
+        return intArrayOf(
+            colors[0],
+            colors[1],
+            colors[0],
+            colors[1],
         )
     }
 
@@ -307,7 +373,7 @@
         }
     }
 
-    private suspend fun loadPreset() =
+    private suspend fun loadPreset(shouldUseRevampedUi: Boolean) =
         withContext(Dispatchers.IO) {
             val extractor = ColorBundlePreviewExtractor(mContext)
             val bundles: MutableList<ColorOption> = ArrayList()
@@ -316,52 +382,134 @@
             // Color option index value starts from 1.
             var index = 1
             val maxPresetColors = if (themeStyleEnabled) bundleNames.size else MAX_PRESET_COLORS
-            for (bundleName in bundleNames.take(maxPresetColors)) {
-                val builder = ColorBundle.Builder()
-                builder.title = getItemStringFromStub(COLOR_BUNDLE_NAME_PREFIX, bundleName)
-                builder.setIndex(index)
-                val colorFromStub = getItemColorFromStub(COLOR_BUNDLE_MAIN_COLOR_PREFIX, bundleName)
-                extractor.addPrimaryColor(builder, colorFromStub)
-                extractor.addSecondaryColor(builder, colorFromStub)
-                if (themeStyleEnabled) {
-                    val styleName =
-                        try {
-                            getItemStringFromStub(COLOR_BUNDLE_STYLE_PREFIX, bundleName)
-                        } catch (e: Resources.NotFoundException) {
-                            null
-                        }
-                    extractor.addColorStyle(builder, styleName)
-                    val style =
-                        try {
-                            if (styleName != null) Style.valueOf(styleName) else Style.TONAL_SPOT
-                        } catch (e: IllegalArgumentException) {
-                            Style.TONAL_SPOT
-                        }
 
-                    if (
-                        style == Style.MONOCHROMATIC &&
-                            !InjectorProvider.getInjector()
-                                .getFlags()
-                                .isMonochromaticThemeEnabled(mContext)
-                    ) {
-                        continue
+            if (shouldUseRevampedUi) {
+                // keep track of whether monochrome is included in preset colors to determine
+                // inclusion in wallpaper colors
+                var hasMonochrome = false
+                for (bundleName in bundleNames.take(maxPresetColors)) {
+                    if (themeStyleEnabled) {
+                        val styleName =
+                            try {
+                                getItemStringFromStub(COLOR_BUNDLE_STYLE_PREFIX, bundleName)
+                            } catch (e: Resources.NotFoundException) {
+                                null
+                            }
+                        val style =
+                            try {
+                                if (styleName != null) Style.valueOf(styleName)
+                                else Style.TONAL_SPOT
+                            } catch (e: IllegalArgumentException) {
+                                Style.TONAL_SPOT
+                            }
+
+                        if (style == Style.MONOCHROMATIC) {
+                            if (!monochromeEnabled) {
+                                continue
+                            }
+                            hasMonochrome = true
+                            monochromeBundleName = bundleName
+                        }
+                        bundles.add(buildRevampedUIPreset(bundleName, index, style))
+                    } else {
+                        bundles.add(buildRevampedUIPreset(bundleName, index, null))
                     }
 
-                    val darkColors =
-                        ColorScheme(colorFromStub, true, style).getPresetColorPreview(colorFromStub)
-                    val lightColors =
-                        ColorScheme(colorFromStub, false, style)
-                            .getPresetColorPreview(colorFromStub)
-                    builder.setColorPrimaryDark(darkColors[0]).setColorSecondaryDark(darkColors[1])
-                    builder
-                        .setColorPrimaryLight(lightColors[0])
-                        .setColorSecondaryLight(lightColors[1])
+                    index++
                 }
+                if (!hasMonochrome) {
+                    monochromeBundleName = null
+                }
+            } else {
+                for (bundleName in bundleNames.take(maxPresetColors)) {
+                    val builder = ColorBundle.Builder()
+                    builder.title = getItemStringFromStub(COLOR_BUNDLE_NAME_PREFIX, bundleName)
+                    builder.setIndex(index)
+                    val colorFromStub =
+                        getItemColorFromStub(COLOR_BUNDLE_MAIN_COLOR_PREFIX, bundleName)
+                    extractor.addPrimaryColor(builder, colorFromStub)
+                    extractor.addSecondaryColor(builder, colorFromStub)
+                    if (themeStyleEnabled) {
+                        val styleName =
+                            try {
+                                getItemStringFromStub(COLOR_BUNDLE_STYLE_PREFIX, bundleName)
+                            } catch (e: Resources.NotFoundException) {
+                                null
+                            }
+                        extractor.addColorStyle(builder, styleName)
+                        val style =
+                            try {
+                                if (styleName != null) Style.valueOf(styleName)
+                                else Style.TONAL_SPOT
+                            } catch (e: IllegalArgumentException) {
+                                Style.TONAL_SPOT
+                            }
 
-                bundles.add(builder.build(mContext))
-                index++
+                        if (style == Style.MONOCHROMATIC && !monochromeEnabled) {
+                            continue
+                        }
+
+                        val darkColors =
+                            ColorScheme(colorFromStub, /* darkTheme= */ true, style)
+                                .getPresetColorPreview(colorFromStub)
+                        val lightColors =
+                            ColorScheme(colorFromStub, /* darkTheme= */ false, style)
+                                .getPresetColorPreview(colorFromStub)
+                        builder
+                            .setColorPrimaryDark(darkColors[0])
+                            .setColorSecondaryDark(darkColors[1])
+                        builder
+                            .setColorPrimaryLight(lightColors[0])
+                            .setColorSecondaryLight(lightColors[1])
+                    }
+
+                    bundles.add(builder.build(mContext))
+                    index++
+                }
             }
 
             colorBundles = bundles
         }
+
+    private fun buildRevampedUIPreset(
+        bundleName: String,
+        index: Int,
+        style: Style? = null,
+        type: ColorType = ColorType.PRESET_COLOR,
+    ): ColorOptionImpl {
+        val builder = ColorOptionImpl.Builder()
+        builder.title = getItemStringFromStub(COLOR_BUNDLE_NAME_PREFIX, bundleName)
+        builder.index = index
+        builder.source = ColorOptionsProvider.COLOR_SOURCE_PRESET
+        builder.type = type
+        val colorFromStub = getItemColorFromStub(COLOR_BUNDLE_MAIN_COLOR_PREFIX, bundleName)
+        var darkColorScheme = ColorScheme(colorFromStub, /* darkTheme= */ true)
+        var lightColorScheme = ColorScheme(colorFromStub, /* darkTheme= */ false)
+        val lightColor = lightColorScheme.accentColor
+        val darkColor = darkColorScheme.accentColor
+        var lightColors = intArrayOf(lightColor, lightColor, lightColor, lightColor)
+        var darkColors = intArrayOf(darkColor, darkColor, darkColor, darkColor)
+        builder.addOverlayPackage(OVERLAY_CATEGORY_COLOR, toColorString(colorFromStub))
+        builder.addOverlayPackage(OVERLAY_CATEGORY_SYSTEM_PALETTE, toColorString(colorFromStub))
+        if (style != null) {
+            builder.style = style
+
+            lightColorScheme = ColorScheme(colorFromStub, /* darkTheme= */ false, style)
+            darkColorScheme = ColorScheme(colorFromStub, /* darkTheme= */ true, style)
+
+            when (style) {
+                Style.MONOCHROMATIC -> {
+                    darkColors = getRevampedUIDarkColorPreview(darkColorScheme)
+                    lightColors = getRevampedUILightColorPreview(lightColorScheme)
+                }
+                else -> {
+                    darkColors = getRevampedUIPresetColorPreview(darkColorScheme, colorFromStub)
+                    lightColors = getRevampedUIPresetColorPreview(lightColorScheme, colorFromStub)
+                }
+            }
+        }
+        builder.lightColors = lightColors
+        builder.darkColors = darkColors
+        return builder.build()
+    }
 }
diff --git a/src/com/android/customization/picker/clock/ui/view/ClockCarouselView.kt b/src/com/android/customization/picker/clock/ui/view/ClockCarouselView.kt
index a27d019..a2d5aec 100644
--- a/src/com/android/customization/picker/clock/ui/view/ClockCarouselView.kt
+++ b/src/com/android/customization/picker/clock/ui/view/ClockCarouselView.kt
@@ -98,11 +98,11 @@
                     endId: Int,
                     progress: Float
                 ) {
-                    scalingDownClockController.animations.onPickerCarouselSwiping(
+                    scalingDownClockController.largeClock.animations.onPickerCarouselSwiping(
                         1 - progress,
                         getPreviewRatio()
                     )
-                    scalingUpClockController.animations.onPickerCarouselSwiping(
+                    scalingUpClockController.largeClock.animations.onPickerCarouselSwiping(
                         progress,
                         getPreviewRatio()
                     )
@@ -173,6 +173,7 @@
                 clockScaleView.scaleX = CLOCK_CAROUSEL_VIEW_SCALE
                 clockScaleView.scaleY = CLOCK_CAROUSEL_VIEW_SCALE
                 onGetClockController(clockIds[index])
+                    .largeClock
                     .animations
                     .onPickerCarouselSwiping(0F, getPreviewRatio())
             } else {
@@ -180,6 +181,7 @@
                 clockScaleView.scaleX = 1f
                 clockScaleView.scaleY = 1f
                 onGetClockController(clockIds[index])
+                    .largeClock
                     .animations
                     .onPickerCarouselSwiping(1F, getPreviewRatio())
             }
diff --git a/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModel.kt b/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModel.kt
index a0ac370..b0ff1db 100644
--- a/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModel.kt
+++ b/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModel.kt
@@ -20,8 +20,7 @@
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.viewModelScope
-import com.android.customization.model.color.ColorBundle
-import com.android.customization.model.color.ColorSeedOption
+import com.android.customization.model.color.ColorOptionImpl
 import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
 import com.android.customization.picker.clock.shared.ClockSize
 import com.android.customization.picker.clock.shared.model.ClockMetadataModel
@@ -201,52 +200,32 @@
     private suspend fun ColorOptionModel.toOptionItemViewModel(
         context: Context
     ): OptionItemViewModel<ColorOptionIconViewModel> {
-        val optionItemPayload =
-            when (colorOption) {
-                is ColorSeedOption -> {
-                    val lightThemeColors =
-                        colorOption.previewInfo.resolveColors(
-                            /** darkTheme= */
-                            false
-                        )
-                    val darkThemeColors =
-                        colorOption.previewInfo.resolveColors(
-                            /** darkTheme= */
-                            true
-                        )
-                    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],
-                    )
-                }
-                is ColorBundle -> {
-                    val primaryColor =
-                        colorOption.previewInfo.resolvePrimaryColor(context.resources)
-                    val secondaryColor =
-                        colorOption.previewInfo.resolveSecondaryColor(context.resources)
-                    ColorOptionIconViewModel(
-                        lightThemeColor0 = primaryColor,
-                        lightThemeColor1 = secondaryColor,
-                        lightThemeColor2 = primaryColor,
-                        lightThemeColor3 = secondaryColor,
-                        darkThemeColor0 = primaryColor,
-                        darkThemeColor1 = secondaryColor,
-                        darkThemeColor2 = primaryColor,
-                        darkThemeColor3 = secondaryColor,
-                    )
-                }
-                else -> null
-            }
+        val lightThemeColors =
+            (colorOption as ColorOptionImpl)
+                .previewInfo
+                .resolveColors(
+                    /** darkTheme= */
+                    false
+                )
+        val darkThemeColors =
+            colorOption.previewInfo.resolveColors(
+                /** darkTheme= */
+                true
+            )
         val isSelectedFlow = selectedColorId.map { it == null }.stateIn(viewModelScope)
         return OptionItemViewModel<ColorOptionIconViewModel>(
             key = MutableStateFlow(key) as StateFlow<String>,
-            payload = optionItemPayload,
+            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(context.getString(R.string.default_theme_title)),
             isTextUserVisible = true,
             isSelected = isSelectedFlow,
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 a044174..464e8bf 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
@@ -19,10 +19,9 @@
 import android.app.WallpaperColors
 import android.util.Log
 import com.android.customization.model.CustomizationManager
-import com.android.customization.model.color.ColorBundle
 import com.android.customization.model.color.ColorCustomizationManager
 import com.android.customization.model.color.ColorOption
-import com.android.customization.model.color.ColorSeedOption
+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
@@ -60,10 +59,11 @@
                                 val presetColorOptions: MutableList<ColorOptionModel> =
                                     mutableListOf()
                                 options?.forEach { option ->
-                                    when (option) {
-                                        is ColorSeedOption ->
+                                    when ((option as ColorOptionImpl).type) {
+                                        ColorType.WALLPAPER_COLOR ->
                                             wallpaperColorOptions.add(option.toModel())
-                                        is ColorBundle -> presetColorOptions.add(option.toModel())
+                                        ColorType.PRESET_COLOR ->
+                                            presetColorOptions.add(option.toModel())
                                     }
                                 }
                                 continuation.resumeWith(
@@ -113,16 +113,16 @@
         val overlays = colorManager.currentOverlays
         val styleOrNull = colorManager.currentStyle
         val style = styleOrNull?.let { Style.valueOf(it) } ?: Style.TONAL_SPOT
-        val colorOptionBuilder =
-            // Does not matter whether ColorSeedOption or ColorBundle builder is used here
-            // because to apply the color, one just needs a generic ColorOption
-            ColorSeedOption.Builder().setSource(colorManager.currentColorSource).setStyle(style)
+        val source = colorManager.currentColorSource
+        val colorOptionBuilder = ColorOptionImpl.Builder()
+        colorOptionBuilder.source = source
+        colorOptionBuilder.style = style
         for (overlay in overlays) {
             colorOptionBuilder.addOverlayPackage(overlay.key, overlay.value)
         }
         val colorOption = colorOptionBuilder.build()
         return ColorOptionModel(
-            key = "${colorOption.style}::${colorOption.serializedPackages}",
+            key = "",
             colorOption = colorOption,
             isSelected = false,
         )
@@ -132,9 +132,9 @@
         return colorManager.currentColorSource
     }
 
-    private fun ColorOption.toModel(): ColorOptionModel {
+    private fun ColorOptionImpl.toModel(): ColorOptionModel {
         return ColorOptionModel(
-            key = "${this.style}::${this.serializedPackages}",
+            key = "${this.type}::${this.style}::${this.serializedPackages}",
             colorOption = this,
             isSelected = 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 edbf6dc..714129d 100644
--- a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
@@ -19,9 +19,8 @@
 import android.content.Context
 import android.graphics.Color
 import android.text.TextUtils
-import com.android.customization.model.color.ColorBundle
+import com.android.customization.model.color.ColorOptionImpl
 import com.android.customization.model.color.ColorOptionsProvider
-import com.android.customization.model.color.ColorSeedOption
 import com.android.customization.picker.color.shared.model.ColorOptionModel
 import com.android.customization.picker.color.shared.model.ColorType
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -82,9 +81,7 @@
                                 ColorOptionModel(
                                     key = "${ColorType.PRESET_COLOR}::$index",
                                     colorOption = buildPresetOption(index),
-                                    isSelected =
-                                        selectedColorOptionType == ColorType.PRESET_COLOR &&
-                                            selectedColorOptionIndex == index,
+                                    isSelected = isSelected,
                                 )
                             if (isSelected) {
                                 selectedColorOption = colorOption
@@ -95,36 +92,36 @@
             )
     }
 
-    private fun buildPresetOption(index: Int): ColorBundle {
-        return ColorBundle.Builder()
+    private fun buildPresetOption(index: Int): ColorOptionImpl {
+        val builder = ColorOptionImpl.Builder()
+        builder.lightColors =
+            intArrayOf(Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT)
+        builder.darkColors =
+            intArrayOf(Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT)
+        builder.index = index
+        builder.type = ColorType.PRESET_COLOR
+        builder.source = ColorOptionsProvider.COLOR_SOURCE_PRESET
+        builder.title = "Preset"
+        builder
             .addOverlayPackage("TEST_PACKAGE_TYPE", "preset_color")
             .addOverlayPackage("TEST_PACKAGE_INDEX", "$index")
-            .setIndex(index)
-            .build(context)
+        return builder.build()
     }
 
-    private fun buildWallpaperOption(index: Int): ColorSeedOption {
-        return ColorSeedOption.Builder()
-            .setLightColors(
-                intArrayOf(
-                    Color.TRANSPARENT,
-                    Color.TRANSPARENT,
-                    Color.TRANSPARENT,
-                    Color.TRANSPARENT
-                )
-            )
-            .setDarkColors(
-                intArrayOf(
-                    Color.TRANSPARENT,
-                    Color.TRANSPARENT,
-                    Color.TRANSPARENT,
-                    Color.TRANSPARENT
-                )
-            )
+    private fun buildWallpaperOption(index: Int): ColorOptionImpl {
+        val builder = ColorOptionImpl.Builder()
+        builder.lightColors =
+            intArrayOf(Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT)
+        builder.darkColors =
+            intArrayOf(Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT)
+        builder.index = index
+        builder.type = ColorType.WALLPAPER_COLOR
+        builder.source = ColorOptionsProvider.COLOR_SOURCE_HOME
+        builder.title = "Dynamic"
+        builder
             .addOverlayPackage("TEST_PACKAGE_TYPE", "wallpaper_color")
             .addOverlayPackage("TEST_PACKAGE_INDEX", "$index")
-            .setIndex(index)
-            .build()
+        return builder.build()
     }
 
     override suspend fun select(colorOptionModel: ColorOptionModel) {
@@ -163,9 +160,9 @@
     override fun getCurrentColorOption(): ColorOptionModel = selectedColorOption
 
     override fun getCurrentColorSource(): String? =
-        when (selectedColorOption.colorOption) {
-            is ColorSeedOption -> ColorOptionsProvider.COLOR_SOURCE_HOME
-            is ColorBundle -> ColorOptionsProvider.COLOR_SOURCE_PRESET
+        when ((selectedColorOption.colorOption as ColorOptionImpl).type) {
+            ColorType.WALLPAPER_COLOR -> ColorOptionsProvider.COLOR_SOURCE_HOME
+            ColorType.PRESET_COLOR -> ColorOptionsProvider.COLOR_SOURCE_PRESET
             else -> null
         }
 
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 97dbba6..119584d 100644
--- a/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
+++ b/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
@@ -20,8 +20,7 @@
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.viewModelScope
-import com.android.customization.model.color.ColorBundle
-import com.android.customization.model.color.ColorSeedOption
+import com.android.customization.model.color.ColorOptionImpl
 import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
 import com.android.customization.picker.color.shared.model.ColorType
 import com.android.wallpaper.R
@@ -96,125 +95,52 @@
             colorOptions
                 .map { colorOptionEntry ->
                     colorOptionEntry.key to
-                        when (colorOptionEntry.key) {
-                            ColorType.WALLPAPER_COLOR -> {
-                                colorOptionEntry.value.map { colorOptionModel ->
-                                    val colorSeedOption: ColorSeedOption =
-                                        colorOptionModel.colorOption as ColorSeedOption
-                                    val lightThemeColors =
-                                        colorSeedOption.previewInfo.resolveColors(
-                                            /* darkTheme= */ false
-                                        )
-                                    val darkThemeColors =
-                                        colorSeedOption.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(
-                                                colorSeedOption
-                                                    .getContentDescription(context)
-                                                    .toString()
-                                            ),
-                                        isTextUserVisible = false,
-                                        isSelected = isSelectedFlow,
-                                        onClicked =
-                                            isSelectedFlow.map { isSelected ->
-                                                if (isSelected) {
-                                                    null
-                                                } else {
-                                                    {
-                                                        viewModelScope.launch {
-                                                            interactor.select(colorOptionModel)
-                                                        }
-                                                    }
+                        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)
                                                 }
-                                            },
-                                    )
-                                }
-                            }
-                            ColorType.PRESET_COLOR -> {
-                                colorOptionEntry.value.map { colorOptionModel ->
-                                    val colorBundle: ColorBundle =
-                                        colorOptionModel.colorOption as ColorBundle
-                                    val primaryColor =
-                                        colorBundle.previewInfo.resolvePrimaryColor(
-                                            context.resources
-                                        )
-                                    val secondaryColor =
-                                        colorBundle.previewInfo.resolveSecondaryColor(
-                                            context.resources
-                                        )
-                                    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 = primaryColor,
-                                                lightThemeColor1 = secondaryColor,
-                                                lightThemeColor2 = primaryColor,
-                                                lightThemeColor3 = secondaryColor,
-                                                darkThemeColor0 = primaryColor,
-                                                darkThemeColor1 = secondaryColor,
-                                                darkThemeColor2 = primaryColor,
-                                                darkThemeColor3 = secondaryColor,
-                                            ),
-                                        text =
-                                            Text.Loaded(
-                                                colorBundle
-                                                    .getContentDescription(context)
-                                                    .toString()
-                                            ),
-                                        isTextUserVisible = false,
-                                        isSelected = isSelectedFlow,
-                                        onClicked =
-                                            isSelectedFlow.map { isSelected ->
-                                                if (isSelected) {
-                                                    null
-                                                } else {
-                                                    {
-                                                        viewModelScope.launch {
-                                                            interactor.select(colorOptionModel)
-                                                        }
-                                                    }
-                                                }
-                                            },
-                                    )
-                                }
-                            }
+                                        }
+                                    },
+                            )
                         }
                 }
                 .toMap()
diff --git a/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt b/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
index 1d9457a..a29d559 100644
--- a/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
+++ b/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
@@ -61,8 +61,8 @@
     fun setUp() {
         context = InstrumentationRegistry.getInstrumentation().targetContext
         val testDispatcher = StandardTestDispatcher()
-        testScope = TestScope(testDispatcher)
         Dispatchers.setMain(testDispatcher)
+        testScope = TestScope(testDispatcher)
         repository = FakeColorPickerRepository(context = context)
         store = FakeSnapshotStore()