Merge tag 'android-13.0.0_r52' of https://android.googlesource.com/platform/packages/apps/ThemePicker into HEADHEADt13.0

Android 13.0.0 Release 52 (TQ3A.230605.012)

Change-Id: I2cea11fa2f1f02fbd3c9d21cfc1697a79d42a5b7

27 files changed, 1169 insertions, 140 deletions

diff --git a/src/com/android/customization/model/clock/ClockSectionController.kt b/src/com/android/customization/model/clock/ClockSectionController.kt
deleted file mode 100644
index 1e339bb1..00000000
--- a/src/com/android/customization/model/clock/ClockSectionController.kt
+++ /dev/null
@@ -1,50 +0,0 @@
-/*
- * Copyright (C) 2022 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.clock
-
-import android.content.Context
-import android.view.LayoutInflater
-import com.android.customization.picker.clock.ClockCustomDemoFragment
-import com.android.customization.picker.clock.ClockSectionView
-import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
-import com.android.systemui.shared.customization.data.content.CustomizationProviderContract as Contract
-import com.android.wallpaper.R
-import com.android.wallpaper.model.CustomizationSectionController
-import com.android.wallpaper.model.CustomizationSectionController.CustomizationSectionNavigationController
-import kotlinx.coroutines.runBlocking
-
-/** A [CustomizationSectionController] for clock customization. */
-class ClockSectionController(
-    private val navigationController: CustomizationSectionNavigationController,
-    private val customizationProviderClient: CustomizationProviderClient,
-) : CustomizationSectionController<ClockSectionView?> {
-    override fun isAvailable(context: Context?): Boolean {
-        return runBlocking { customizationProviderClient.queryFlags() }
-            .firstOrNull { it.name == Contract.FlagsTable.FLAG_NAME_CUSTOM_CLOCKS_ENABLED }
-            ?.value == true
-    }
-
-    override fun createView(context: Context): ClockSectionView {
-        val view =
-            LayoutInflater.from(context)
-                .inflate(
-                    R.layout.clock_section_view,
-                    null,
-                ) as ClockSectionView
-        view.setOnClickListener { navigationController.navigateTo(ClockCustomDemoFragment()) }
-        return view
-    }
-}
diff --git a/src/com/android/customization/model/color/ColorCustomizationManager.java b/src/com/android/customization/model/color/ColorCustomizationManager.java
index 908480f6..29f6ba6e 100644
--- a/src/com/android/customization/model/color/ColorCustomizationManager.java
+++ b/src/com/android/customization/model/color/ColorCustomizationManager.java
@@ -210,7 +210,7 @@ public class ColorCustomizationManager implements CustomizationManager<ColorOpti
      * or {@link ColorOptionsProvider#COLOR_SOURCE_PRESET}.
      */
     @ColorSource
-    public String getCurrentColorSource() {
+    public @Nullable String getCurrentColorSource() {
         if (mCurrentSource == null) {
             parseSettings(getStoredOverlays());
         }
@@ -221,7 +221,7 @@ public class ColorCustomizationManager implements CustomizationManager<ColorOpti
      * @return The style of the currently applied color. One of enum values in
      * {@link com.android.systemui.monet.Style}.
      */
-    public String getCurrentStyle() {
+    public @Nullable String getCurrentStyle() {
         if (mCurrentStyle == null) {
             parseSettings(getStoredOverlays());
         }
diff --git a/src/com/android/customization/model/color/ColorOption.java b/src/com/android/customization/model/color/ColorOption.java
index c8b28c29..216bb9ba 100644
--- a/src/com/android/customization/model/color/ColorOption.java
+++ b/src/com/android/customization/model/color/ColorOption.java
@@ -107,9 +107,15 @@ public abstract class ColorOption implements CustomizationOption<ColorOption> {
         if (other == null) {
             return false;
         }
-        if (mIsDefault) {
-            return other.isDefault() || TextUtils.isEmpty(other.getSerializedPackages())
-                    || EMPTY_JSON.equals(other.getSerializedPackages());
+        if (mStyle != other.getStyle()) {
+            return false;
+        }
+        String thisSerializedPackages = getSerializedPackages();
+        if (mIsDefault || TextUtils.isEmpty(thisSerializedPackages)
+                || EMPTY_JSON.equals(thisSerializedPackages)) {
+            String otherSerializedPackages = other.getSerializedPackages();
+            return other.isDefault() || TextUtils.isEmpty(otherSerializedPackages)
+                    || EMPTY_JSON.equals(otherSerializedPackages);
         }
         // Map#equals ensures keys and values are compared.
         return mPackagesByCategory.equals(other.mPackagesByCategory);
@@ -182,7 +188,8 @@ public abstract class ColorOption implements CustomizationOption<ColorOption> {
                 .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
     }
 
-    protected CharSequence getContentDescription(Context context) {
+    /** */
+    public CharSequence getContentDescription(Context context) {
         if (mContentDescription == null) {
             CharSequence defaultName = context.getString(R.string.default_theme_title);
             if (isDefault()) {
diff --git a/src/com/android/customization/model/color/ColorProvider.kt b/src/com/android/customization/model/color/ColorProvider.kt
index 3d2cc7ef..f91ec6b7 100644
--- a/src/com/android/customization/model/color/ColorProvider.kt
+++ b/src/com/android/customization/model/color/ColorProvider.kt
@@ -220,12 +220,22 @@ class ColorProvider(context: Context, stubPackageName: String) :
      */
     @ColorInt
     private fun ColorScheme.getLightColorPreview(): IntArray {
-        return intArrayOf(
-            setAlphaComponent(this.accent1[2], ALPHA_MASK),
-            setAlphaComponent(this.accent1[2], ALPHA_MASK),
-            ColorStateList.valueOf(this.accent3[6]).withLStar(85f).colors[0],
-            setAlphaComponent(this.accent1[6], ALPHA_MASK)
-        )
+        return when (this.style) {
+            Style.EXPRESSIVE ->
+                intArrayOf(
+                    setAlphaComponent(this.accent1.s100, ALPHA_MASK),
+                    setAlphaComponent(this.accent1.s100, ALPHA_MASK),
+                    ColorStateList.valueOf(this.neutral2.s500).withLStar(80f).colors[0],
+                    setAlphaComponent(this.accent2.s500, ALPHA_MASK)
+                )
+            else ->
+                intArrayOf(
+                    setAlphaComponent(this.accent1.s100, ALPHA_MASK),
+                    setAlphaComponent(this.accent1.s100, ALPHA_MASK),
+                    ColorStateList.valueOf(this.accent3.s500).withLStar(85f).colors[0],
+                    setAlphaComponent(this.accent1.s500, ALPHA_MASK)
+                )
+        }
     }
 
     /**
@@ -234,24 +244,19 @@ class ColorProvider(context: Context, stubPackageName: String) :
      */
     @ColorInt
     private fun ColorScheme.getDarkColorPreview(): IntArray {
-        return intArrayOf(
-            setAlphaComponent(this.accent1[2], ALPHA_MASK),
-            setAlphaComponent(this.accent1[2], ALPHA_MASK),
-            ColorStateList.valueOf(this.accent3[6]).withLStar(85f).colors[0],
-            setAlphaComponent(this.accent1[6], ALPHA_MASK)
-        )
+        return getLightColorPreview()
     }
 
     private fun ColorScheme.getPresetColorPreview(seed: Int): IntArray {
         return when (this.style) {
-            Style.FRUIT_SALAD -> intArrayOf(seed, this.accent1[2])
+            Style.FRUIT_SALAD -> intArrayOf(seed, this.accent1.s100)
             Style.TONAL_SPOT -> intArrayOf(this.accentColor, this.accentColor)
             Style.MONOCHROMATIC ->
                 intArrayOf(
                     setAlphaComponent(0x000000, 255),
                     setAlphaComponent(0xFFFFFF, 255),
                 )
-            else -> intArrayOf(this.accent1[2], this.accent1[2])
+            else -> intArrayOf(this.accent1.s100, this.accent1.s100)
         }
     }
 
@@ -288,7 +293,9 @@ class ColorProvider(context: Context, stubPackageName: String) :
 
                     if (
                         style == Style.MONOCHROMATIC &&
-                            !InjectorProvider.getInjector().getFlags().isMonochromaticFlagEnabled()
+                            !InjectorProvider.getInjector()
+                                .getFlags()
+                                .isMonochromaticThemeEnabled(mContext)
                     ) {
                         continue
                     }
diff --git a/src/com/android/customization/model/color/ColorSectionController.java b/src/com/android/customization/model/color/ColorSectionController.java
index 3b8a9273..be051ac0 100644
--- a/src/com/android/customization/model/color/ColorSectionController.java
+++ b/src/com/android/customization/model/color/ColorSectionController.java
@@ -41,7 +41,6 @@ import android.widget.FrameLayout;
 import androidx.annotation.Nullable;
 import androidx.lifecycle.LifecycleOwner;
 import androidx.recyclerview.widget.RecyclerView;
-import androidx.viewpager2.widget.MarginPageTransformer;
 import androidx.viewpager2.widget.ViewPager2;
 
 import com.android.customization.model.CustomizationManager;
@@ -54,7 +53,6 @@ import com.android.wallpaper.R;
 import com.android.wallpaper.model.CustomizationSectionController;
 import com.android.wallpaper.model.WallpaperColorsViewModel;
 import com.android.wallpaper.module.InjectorProvider;
-import com.android.wallpaper.module.LargeScreenMultiPanesChecker;
 import com.android.wallpaper.widget.PageIndicator;
 import com.android.wallpaper.widget.SeparatedTabLayout;
 
@@ -104,7 +102,6 @@ public class ColorSectionController implements CustomizationSectionController<Co
             new Optional[]{Optional.empty(), Optional.empty()};
     private long mLastColorApplyingTime = 0L;
     private ColorSectionView mColorSectionView;
-    private boolean mIsMultiPane;
 
     private static int getNumPages(int optionsPerPage, int totalOptions) {
         return (int) Math.ceil((float) totalOptions / optionsPerPage);
@@ -118,7 +115,6 @@ public class ColorSectionController implements CustomizationSectionController<Co
                 new OverlayManagerCompat(activity));
         mWallpaperColorsViewModel = viewModel;
         mLifecycleOwner = lifecycleOwner;
-        mIsMultiPane = new LargeScreenMultiPanesChecker().isMultiPanesEnabled(activity);
 
         if (savedInstanceState != null) {
             if (savedInstanceState.containsKey(KEY_COLOR_TAB_POSITION)) {
@@ -174,13 +170,13 @@ public class ColorSectionController implements CustomizationSectionController<Co
         // TODO(b/202145216): Use just 2 views when tapping either button on top.
         mTabLayout.setViewPager(mColorSectionViewPager);
 
-        mWallpaperColorsViewModel.getHomeWallpaperColors().observe(mLifecycleOwner,
+        mWallpaperColorsViewModel.getHomeWallpaperColorsLiveData().observe(mLifecycleOwner,
                 homeColors -> {
                     mHomeWallpaperColors = homeColors;
                     mHomeWallpaperColorsReady = true;
                     maybeLoadColors();
                 });
-        mWallpaperColorsViewModel.getLockWallpaperColors().observe(mLifecycleOwner,
+        mWallpaperColorsViewModel.getLockWallpaperColorsLiveData().observe(mLifecycleOwner,
                 lockColors -> {
                     mLockWallpaperColors = lockColors;
                     mLockWallpaperColorsReady = true;
@@ -472,16 +468,6 @@ public class ColorSectionController implements CustomizationSectionController<Co
                 mContainer = itemView.findViewById(R.id.color_page_container);
                 // Correct scrolling goes under collapsing toolbar while scrolling oclor options.
                 mContainer.getChildAt(0).setNestedScrollingEnabled(false);
-                /**
-                 * Sets page transformer with margin to separate color pages and
-                 * sets color pages' padding to not scroll to window boundary if multi-pane case
-                 */
-                if (mIsMultiPane) {
-                    final int padding = itemView.getContext().getResources().getDimensionPixelSize(
-                            R.dimen.section_horizontal_padding);
-                    mContainer.setPageTransformer(new MarginPageTransformer(padding * 2));
-                    mContainer.setPadding(padding, /* top= */ 0, padding, /* bottom= */ 0);
-                }
                 mPageIndicator = itemView.findViewById(R.id.color_page_indicator);
                 if (ColorProvider.themeStyleEnabled) {
                     mPageIndicator.setVisibility(VISIBLE);
@@ -539,15 +525,13 @@ public class ColorSectionController implements CustomizationSectionController<Co
             ColorOptionViewHolder(View itemView) {
                 super(itemView);
                 mContainer = itemView.findViewById(R.id.color_option_container);
-                // Sets layout with margins for non multi-pane case to separate color options.
-                if (!mIsMultiPane) {
-                    final FrameLayout.LayoutParams layoutParams = new FrameLayout.LayoutParams(
-                            mContainer.getLayoutParams());
-                    final int margin = itemView.getContext().getResources().getDimensionPixelSize(
-                            R.dimen.section_horizontal_padding);
-                    layoutParams.setMargins(margin, /* top= */ 0, margin, /* bottom= */ 0);
-                    mContainer.setLayoutParams(layoutParams);
-                }
+                // Sets layout with margins to separate color options.
+                final FrameLayout.LayoutParams layoutParams = new FrameLayout.LayoutParams(
+                        mContainer.getLayoutParams());
+                final int margin = itemView.getContext().getResources().getDimensionPixelSize(
+                        R.dimen.section_horizontal_padding);
+                layoutParams.setMargins(margin, /* top= */ 0, margin, /* bottom= */ 0);
+                mContainer.setLayoutParams(layoutParams);
             }
         }
     }
diff --git a/src/com/android/customization/model/color/ColorSeedOption.java b/src/com/android/customization/model/color/ColorSeedOption.java
index 53d39543..ba61ed1b 100644
--- a/src/com/android/customization/model/color/ColorSeedOption.java
+++ b/src/com/android/customization/model/color/ColorSeedOption.java
@@ -80,7 +80,7 @@ public class ColorSeedOption extends ColorOption {
     }
 
     @Override
-    protected CharSequence getContentDescription(Context context) {
+    public CharSequence getContentDescription(Context context) {
         // Override because we want all options with the same description.
         return context.getString(R.string.wallpaper_color_title);
     }
diff --git a/src/com/android/customization/model/color/WallpaperColorResources.java b/src/com/android/customization/model/color/WallpaperColorResources.java
index eb8b39be..dc3b9033 100644
--- a/src/com/android/customization/model/color/WallpaperColorResources.java
+++ b/src/com/android/customization/model/color/WallpaperColorResources.java
@@ -21,8 +21,7 @@ import android.util.SparseIntArray;
 import android.widget.RemoteViews.ColorResources;
 
 import com.android.systemui.monet.ColorScheme;
-
-import java.util.List;
+import com.android.systemui.monet.TonalPalette;
 
 /** A class to override colors in a {@link Context} with wallpaper colors. */
 public class WallpaperColorResources {
@@ -43,9 +42,9 @@ public class WallpaperColorResources {
         ColorResources.create(context, mColorOverlay).apply(context);
     }
 
-    private void addOverlayColor(List<Integer> colors, int firstResourceColorId) {
+    private void addOverlayColor(TonalPalette colorSchemehue, int firstResourceColorId) {
         int resourceColorId = firstResourceColorId;
-        for (int color : colors) {
+        for (int color : colorSchemehue.getAllShades()) {
             mColorOverlay.put(resourceColorId, color);
             resourceColorId++;
         }
diff --git a/src/com/android/customization/model/grid/GridOptionsManager.java b/src/com/android/customization/model/grid/GridOptionsManager.java
index 7f15d836..b7ee37fd 100644
--- a/src/com/android/customization/model/grid/GridOptionsManager.java
+++ b/src/com/android/customization/model/grid/GridOptionsManager.java
@@ -21,7 +21,9 @@ import android.os.Handler;
 import android.os.Looper;
 import android.util.Log;
 
+import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
+import androidx.lifecycle.LiveData;
 
 import com.android.customization.model.CustomizationManager;
 import com.android.customization.module.CustomizationInjector;
@@ -47,6 +49,7 @@ public class GridOptionsManager implements CustomizationManager<GridOption> {
 
     private final LauncherGridOptionsProvider mProvider;
     private final ThemesUserEventLogger mEventLogger;
+    private int mGridOptionSize = -1;
 
     /** Returns the {@link GridOptionsManager} instance. */
     public static GridOptionsManager getInstance(Context context) {
@@ -71,16 +74,17 @@ public class GridOptionsManager implements CustomizationManager<GridOption> {
 
     @Override
     public boolean isAvailable() {
-        int gridOptionSize = 0;
-        try {
-            gridOptionSize = sExecutorService.submit(() -> {
-                List<GridOption> gridOptions = mProvider.fetch(/* reload= */true);
-                return gridOptions == null ? 0 : gridOptions.size();
-            }).get();
-        } catch (InterruptedException | ExecutionException e) {
-            Log.w(TAG, "could not get gridOptionSize", e);
+        if (mGridOptionSize < 0) {
+            try {
+                mGridOptionSize = sExecutorService.submit(() -> {
+                    List<GridOption> gridOptions = mProvider.fetch(/* reload= */true);
+                    return gridOptions == null ? 0 : gridOptions.size();
+                }).get();
+            } catch (InterruptedException | ExecutionException e) {
+                Log.w(TAG, "could not get gridOptionSize", e);
+            }
         }
-        return gridOptionSize > 1 && mProvider.areGridsAvailable();
+        return mGridOptionSize > 1 && mProvider.areGridsAvailable();
     }
 
     @Override
@@ -110,6 +114,13 @@ public class GridOptionsManager implements CustomizationManager<GridOption> {
         });
     }
 
+    /**
+     * Returns an observable that receives a new value each time that the grid options are changed.
+     */
+    public LiveData<Object> getOptionChangeObservable(@Nullable Handler handler) {
+        return mProvider.getOptionChangeObservable(handler);
+    }
+
     /** Call through content provider API to render preview */
     public void renderPreview(Bundle bundle, String gridName,
             PreviewUtils.WorkspacePreviewCallback callback) {
diff --git a/src/com/android/customization/model/grid/GridSectionController.java b/src/com/android/customization/model/grid/GridSectionController.java
index 2f54a1bf..c50bfcc2 100644
--- a/src/com/android/customization/model/grid/GridSectionController.java
+++ b/src/com/android/customization/model/grid/GridSectionController.java
@@ -22,8 +22,12 @@ import android.view.View;
 import android.widget.TextView;
 
 import androidx.annotation.Nullable;
+import androidx.fragment.app.Fragment;
+import androidx.lifecycle.LifecycleOwner;
+import androidx.lifecycle.Observer;
 
 import com.android.customization.model.CustomizationManager.OptionsFetchedListener;
+import com.android.customization.model.grid.ui.fragment.GridFragment2;
 import com.android.customization.picker.grid.GridFragment;
 import com.android.customization.picker.grid.GridSectionView;
 import com.android.wallpaper.R;
@@ -38,11 +42,22 @@ public class GridSectionController implements CustomizationSectionController<Gri
 
     private final GridOptionsManager mGridOptionsManager;
     private final CustomizationSectionNavigationController mSectionNavigationController;
+    private final boolean mIsRevampedUiEnabled;
+    private final Observer<Object> mOptionChangeObserver;
+    private final LifecycleOwner mLifecycleOwner;
+    private TextView mSectionDescription;
+    private View mSectionTile;
 
-    public GridSectionController(GridOptionsManager gridOptionsManager,
-            CustomizationSectionNavigationController sectionNavigationController) {
+    public GridSectionController(
+            GridOptionsManager gridOptionsManager,
+            CustomizationSectionNavigationController sectionNavigationController,
+            LifecycleOwner lifecycleOwner,
+            boolean isRevampedUiEnabled) {
         mGridOptionsManager = gridOptionsManager;
         mSectionNavigationController = sectionNavigationController;
+        mIsRevampedUiEnabled = isRevampedUiEnabled;
+        mLifecycleOwner = lifecycleOwner;
+        mOptionChangeObserver = o -> updateUi(/* reload= */ true);
     }
 
     @Override
@@ -52,34 +67,68 @@ public class GridSectionController implements CustomizationSectionController<Gri
 
     @Override
     public GridSectionView createView(Context context) {
-        GridSectionView gridSectionView = (GridSectionView) LayoutInflater.from(context)
+        final GridSectionView gridSectionView = (GridSectionView) LayoutInflater.from(context)
                 .inflate(R.layout.grid_section_view, /* root= */ null);
-        TextView sectionDescription = gridSectionView.findViewById(R.id.grid_section_description);
-        View sectionTile = gridSectionView.findViewById(R.id.grid_section_tile);
+        mSectionDescription = gridSectionView.findViewById(R.id.grid_section_description);
+        mSectionTile = gridSectionView.findViewById(R.id.grid_section_tile);
 
         // Fetch grid options to show currently set grid.
-        mGridOptionsManager.fetchOptions(new OptionsFetchedListener<GridOption>() {
-            @Override
-            public void onOptionsLoaded(List<GridOption> options) {
-                sectionDescription.setText(getActiveOption(options).getTitle());
-            }
-
-            @Override
-            public void onError(@Nullable Throwable throwable) {
-                if (throwable != null) {
-                    Log.e(TAG, "Error loading grid options", throwable);
-                }
-                sectionDescription.setText(R.string.something_went_wrong);
-                sectionTile.setVisibility(View.GONE);
-            }
-        }, /* The result is getting when calling isAvailable(), so reload= */ false);
+        updateUi(/* The result is getting when calling isAvailable(), so reload= */ false);
+        if (mIsRevampedUiEnabled) {
+            mGridOptionsManager.getOptionChangeObservable(/* handler= */ null).observe(
+                    mLifecycleOwner,
+                    mOptionChangeObserver);
+        }
 
         gridSectionView.setOnClickListener(
-                v -> mSectionNavigationController.navigateTo(new GridFragment()));
+                v -> {
+                    final Fragment gridFragment;
+                    if (mIsRevampedUiEnabled) {
+                        gridFragment = new GridFragment2();
+                    } else {
+                        gridFragment = new GridFragment();
+                    }
+                    mSectionNavigationController.navigateTo(gridFragment);
+                });
 
         return gridSectionView;
     }
 
+    @Override
+    public void release() {
+        if (mIsRevampedUiEnabled && mGridOptionsManager.isAvailable()) {
+            mGridOptionsManager.getOptionChangeObservable(/* handler= */ null).removeObserver(
+                    mOptionChangeObserver
+            );
+        }
+    }
+
+    @Override
+    public void onTransitionOut() {
+        CustomizationSectionController.super.onTransitionOut();
+    }
+
+    private void updateUi(final boolean reload) {
+        mGridOptionsManager.fetchOptions(
+                new OptionsFetchedListener<GridOption>() {
+                    @Override
+                    public void onOptionsLoaded(List<GridOption> options) {
+                        final String title = getActiveOption(options).getTitle();
+                        mSectionDescription.setText(title);
+                    }
+
+                    @Override
+                    public void onError(@Nullable Throwable throwable) {
+                        if (throwable != null) {
+                            Log.e(TAG, "Error loading grid options", throwable);
+                        }
+                        mSectionDescription.setText(R.string.something_went_wrong);
+                        mSectionTile.setVisibility(View.GONE);
+                    }
+                },
+                reload);
+    }
+
     private GridOption getActiveOption(List<GridOption> options) {
         return options.stream()
                 .filter(option -> option.isActive(mGridOptionsManager))
diff --git a/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java b/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
index fd403631..4e775c62 100644
--- a/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
+++ b/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
@@ -19,12 +19,17 @@ import android.content.ContentResolver;
 import android.content.ContentValues;
 import android.content.Context;
 import android.content.res.Resources;
+import android.database.ContentObserver;
 import android.database.Cursor;
+import android.net.Uri;
 import android.os.Bundle;
+import android.os.Handler;
 import android.view.SurfaceView;
 
 import androidx.annotation.Nullable;
 import androidx.annotation.WorkerThread;
+import androidx.lifecycle.LiveData;
+import androidx.lifecycle.MutableLiveData;
 
 import com.android.customization.model.ResourceConstants;
 import com.android.wallpaper.R;
@@ -53,6 +58,7 @@ public class LauncherGridOptionsProvider {
     private final Context mContext;
     private final PreviewUtils mPreviewUtils;
     private List<GridOption> mOptions;
+    private OptionChangeLiveData mLiveData;
 
     public LauncherGridOptionsProvider(Context context, String authorityMetadataKey) {
         mPreviewUtils = new PreviewUtils(context, authorityMetadataKey);
@@ -117,4 +123,52 @@ public class LauncherGridOptionsProvider {
         return mContext.getContentResolver().update(mPreviewUtils.getUri(DEFAULT_GRID), values,
                 null, null);
     }
+
+    /**
+     * Returns an observable that receives a new value each time that the grid options are changed.
+     * Do not call if {@link #areGridsAvailable()} returns false
+     */
+    public LiveData<Object> getOptionChangeObservable(
+            @Nullable Handler handler) {
+        if (mLiveData == null) {
+            mLiveData = new OptionChangeLiveData(
+                    mContext, mPreviewUtils.getUri(DEFAULT_GRID), handler);
+        }
+
+        return mLiveData;
+    }
+
+    private static class OptionChangeLiveData extends MutableLiveData<Object> {
+
+        private final ContentResolver mContentResolver;
+        private final Uri mUri;
+        private final ContentObserver mContentObserver;
+
+        OptionChangeLiveData(
+                Context context,
+                Uri uri,
+                @Nullable Handler handler) {
+            mContentResolver = context.getContentResolver();
+            mUri = uri;
+            mContentObserver = new ContentObserver(handler) {
+                @Override
+                public void onChange(boolean selfChange) {
+                    postValue(new Object());
+                }
+            };
+        }
+
+        @Override
+        protected void onActive() {
+            mContentResolver.registerContentObserver(
+                    mUri,
+                    /* notifyForDescendants= */ true,
+                    mContentObserver);
+        }
+
+        @Override
+        protected void onInactive() {
+            mContentResolver.unregisterContentObserver(mContentObserver);
+        }
+    }
 }
diff --git a/src/com/android/customization/model/grid/data/repository/GridRepository.kt b/src/com/android/customization/model/grid/data/repository/GridRepository.kt
new file mode 100644
index 00000000..9a3be0cc
--- /dev/null
+++ b/src/com/android/customization/model/grid/data/repository/GridRepository.kt
@@ -0,0 +1,127 @@
+/*
+ * 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.grid.data.repository
+
+import androidx.lifecycle.asFlow
+import com.android.customization.model.CustomizationManager
+import com.android.customization.model.grid.GridOption
+import com.android.customization.model.grid.GridOptionsManager
+import com.android.customization.model.grid.shared.model.GridOptionItemModel
+import com.android.customization.model.grid.shared.model.GridOptionItemsModel
+import kotlin.coroutines.resume
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.suspendCancellableCoroutine
+import kotlinx.coroutines.withContext
+
+interface GridRepository {
+    suspend fun isAvailable(): Boolean
+    fun getOptionChanges(): Flow<Unit>
+    suspend fun getOptions(): GridOptionItemsModel
+}
+
+class GridRepositoryImpl(
+    private val applicationScope: CoroutineScope,
+    private val manager: GridOptionsManager,
+    private val backgroundDispatcher: CoroutineDispatcher,
+) : GridRepository {
+
+    override suspend fun isAvailable(): Boolean {
+        return withContext(backgroundDispatcher) { manager.isAvailable }
+    }
+
+    override fun getOptionChanges(): Flow<Unit> =
+        manager.getOptionChangeObservable(/* handler= */ null).asFlow().map {}
+
+    private val selectedOption = MutableStateFlow<GridOption?>(null)
+
+    override suspend fun getOptions(): GridOptionItemsModel {
+        return withContext(backgroundDispatcher) {
+            suspendCancellableCoroutine { continuation ->
+                manager.fetchOptions(
+                    object : CustomizationManager.OptionsFetchedListener<GridOption> {
+                        override fun onOptionsLoaded(options: MutableList<GridOption>?) {
+                            val optionsOrEmpty = options ?: emptyList()
+                            selectedOption.value = optionsOrEmpty.find { it.isActive(manager) }
+                            continuation.resume(
+                                GridOptionItemsModel.Loaded(
+                                    optionsOrEmpty.map { option -> toModel(option) }
+                                )
+                            )
+                        }
+
+                        override fun onError(throwable: Throwable?) {
+                            continuation.resume(
+                                GridOptionItemsModel.Error(
+                                    throwable ?: Exception("Failed to load grid options!")
+                                ),
+                            )
+                        }
+                    },
+                    /* reload= */ true,
+                )
+            }
+        }
+    }
+
+    private fun toModel(option: GridOption): GridOptionItemModel {
+        return GridOptionItemModel(
+            name = option.title,
+            rows = option.rows,
+            cols = option.cols,
+            isSelected =
+                selectedOption
+                    .map { it.key() }
+                    .map { selectedOptionKey -> option.key() == selectedOptionKey }
+                    .stateIn(
+                        scope = applicationScope,
+                        started = SharingStarted.Eagerly,
+                        initialValue = false,
+                    ),
+            onSelected = { onSelected(option) },
+        )
+    }
+
+    private suspend fun onSelected(option: GridOption) {
+        withContext(backgroundDispatcher) {
+            suspendCancellableCoroutine { continuation ->
+                manager.apply(
+                    option,
+                    object : CustomizationManager.Callback {
+                        override fun onSuccess() {
+                            continuation.resume(true)
+                        }
+
+                        override fun onError(throwable: Throwable?) {
+                            continuation.resume(false)
+                        }
+                    },
+                )
+            }
+        }
+    }
+
+    private fun GridOption?.key(): String? {
+        return if (this != null) "${cols}x${rows}" else null
+    }
+}
diff --git a/src/com/android/customization/model/grid/domain/interactor/GridInteractor.kt b/src/com/android/customization/model/grid/domain/interactor/GridInteractor.kt
new file mode 100644
index 00000000..cdb679dd
--- /dev/null
+++ b/src/com/android/customization/model/grid/domain/interactor/GridInteractor.kt
@@ -0,0 +1,98 @@
+/*
+ * 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.grid.domain.interactor
+
+import com.android.customization.model.grid.data.repository.GridRepository
+import com.android.customization.model.grid.shared.model.GridOptionItemModel
+import com.android.customization.model.grid.shared.model.GridOptionItemsModel
+import javax.inject.Provider
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.emptyFlow
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onStart
+import kotlinx.coroutines.flow.shareIn
+
+class GridInteractor(
+    private val applicationScope: CoroutineScope,
+    private val repository: GridRepository,
+    private val snapshotRestorer: Provider<GridSnapshotRestorer>,
+) {
+    val options: Flow<GridOptionItemsModel> =
+        flow { emit(repository.isAvailable()) }
+            .flatMapLatest { isAvailable ->
+                if (isAvailable) {
+                    // this upstream flow tells us each time the options are changed.
+                    repository
+                        .getOptionChanges()
+                        // when we start, we pretend the options _just_ changed. This way, we load
+                        // something as soon as possible into the flow so it's ready by the time the
+                        // first observer starts to observe.
+                        .onStart { emit(Unit) }
+                        // each time the options changed, we load them.
+                        .map { reload() }
+                        // we place the loaded options in a SharedFlow so downstream observers all
+                        // share the same flow and don't trigger a new one each time they want to
+                        // start observing.
+                        .shareIn(
+                            scope = applicationScope,
+                            started = SharingStarted.WhileSubscribed(),
+                            replay = 1,
+                        )
+                } else {
+                    emptyFlow()
+                }
+            }
+
+    suspend fun setSelectedOption(model: GridOptionItemModel) {
+        model.onSelected.invoke()
+    }
+
+    suspend fun getSelectedOption(): GridOptionItemModel? {
+        return (repository.getOptions() as? GridOptionItemsModel.Loaded)?.options?.firstOrNull {
+            optionItem ->
+            optionItem.isSelected.value
+        }
+    }
+
+    private suspend fun reload(): GridOptionItemsModel {
+        val model = repository.getOptions()
+        return if (model is GridOptionItemsModel.Loaded) {
+            GridOptionItemsModel.Loaded(
+                options =
+                    model.options.map { option ->
+                        GridOptionItemModel(
+                            name = option.name,
+                            cols = option.cols,
+                            rows = option.rows,
+                            isSelected = option.isSelected,
+                            onSelected = {
+                                option.onSelected()
+                                snapshotRestorer.get().store(option)
+                            },
+                        )
+                    }
+            )
+        } else {
+            model
+        }
+    }
+}
diff --git a/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorer.kt b/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorer.kt
new file mode 100644
index 00000000..19d4c77e
--- /dev/null
+++ b/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorer.kt
@@ -0,0 +1,74 @@
+/*
+ * 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.grid.domain.interactor
+
+import android.util.Log
+import com.android.customization.model.grid.shared.model.GridOptionItemModel
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotStore
+import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
+
+class GridSnapshotRestorer(
+    private val interactor: GridInteractor,
+) : SnapshotRestorer {
+
+    private var store: SnapshotStore = SnapshotStore.NOOP
+    private var originalOption: GridOptionItemModel? = null
+
+    override suspend fun setUpSnapshotRestorer(store: SnapshotStore): RestorableSnapshot {
+        this.store = store
+        val option = interactor.getSelectedOption()
+        originalOption = option
+        return snapshot(option)
+    }
+
+    override suspend fun restoreToSnapshot(snapshot: RestorableSnapshot) {
+        val optionNameFromSnapshot = snapshot.args[KEY_GRID_OPTION_NAME]
+        originalOption?.let { optionToRestore ->
+            if (optionToRestore.name != optionNameFromSnapshot) {
+                Log.wtf(
+                    TAG,
+                    """Original snapshot name was ${optionToRestore.name} but we're being told to
+                        | restore to $optionNameFromSnapshot. The current implementation doesn't
+                        | support undo, only a reset back to the original grid option."""
+                        .trimMargin(),
+                )
+            }
+
+            interactor.setSelectedOption(optionToRestore)
+        }
+    }
+
+    fun store(option: GridOptionItemModel) {
+        store.store(snapshot(option))
+    }
+
+    private fun snapshot(option: GridOptionItemModel?): RestorableSnapshot {
+        return RestorableSnapshot(
+            args =
+                buildMap {
+                    option?.name?.let { optionName -> put(KEY_GRID_OPTION_NAME, optionName) }
+                }
+        )
+    }
+
+    companion object {
+        private const val TAG = "GridSnapshotRestorer"
+        private const val KEY_GRID_OPTION_NAME = "grid_option"
+    }
+}
diff --git a/src/com/android/customization/model/grid/shared/model/GridOptionItemModel.kt b/src/com/android/customization/model/grid/shared/model/GridOptionItemModel.kt
new file mode 100644
index 00000000..2eabeab5
--- /dev/null
+++ b/src/com/android/customization/model/grid/shared/model/GridOptionItemModel.kt
@@ -0,0 +1,28 @@
+/*
+ * 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.grid.shared.model
+
+import kotlinx.coroutines.flow.StateFlow
+
+data class GridOptionItemModel(
+    val name: String,
+    val cols: Int,
+    val rows: Int,
+    val isSelected: StateFlow<Boolean>,
+    val onSelected: suspend () -> Unit,
+)
diff --git a/src/com/android/customization/model/grid/shared/model/GridOptionItemsModel.kt b/src/com/android/customization/model/grid/shared/model/GridOptionItemsModel.kt
new file mode 100644
index 00000000..e969be88
--- /dev/null
+++ b/src/com/android/customization/model/grid/shared/model/GridOptionItemsModel.kt
@@ -0,0 +1,27 @@
+/*
+ * 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.grid.shared.model
+
+sealed class GridOptionItemsModel {
+    data class Loaded(
+        val options: List<GridOptionItemModel>,
+    ) : GridOptionItemsModel()
+    data class Error(
+        val throwable: Throwable?,
+    ) : GridOptionItemsModel()
+}
diff --git a/src/com/android/customization/model/grid/ui/binder/GridIconViewBinder.kt b/src/com/android/customization/model/grid/ui/binder/GridIconViewBinder.kt
new file mode 100644
index 00000000..fba89a74
--- /dev/null
+++ b/src/com/android/customization/model/grid/ui/binder/GridIconViewBinder.kt
@@ -0,0 +1,17 @@
+package com.android.customization.model.grid.ui.binder
+
+import android.widget.ImageView
+import com.android.customization.model.grid.ui.viewmodel.GridIconViewModel
+import com.android.customization.widget.GridTileDrawable
+
+object GridIconViewBinder {
+    fun bind(view: ImageView, viewModel: GridIconViewModel) {
+        view.setImageDrawable(
+            GridTileDrawable(
+                viewModel.columns,
+                viewModel.rows,
+                viewModel.path,
+            )
+        )
+    }
+}
diff --git a/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt b/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt
new file mode 100644
index 00000000..78536ca9
--- /dev/null
+++ b/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt
@@ -0,0 +1,81 @@
+/*
+ * 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.grid.ui.binder
+
+import android.view.View
+import android.widget.ImageView
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleOwner
+import androidx.lifecycle.lifecycleScope
+import androidx.lifecycle.repeatOnLifecycle
+import androidx.recyclerview.widget.LinearLayoutManager
+import androidx.recyclerview.widget.RecyclerView
+import com.android.customization.model.grid.ui.viewmodel.GridIconViewModel
+import com.android.customization.model.grid.ui.viewmodel.GridScreenViewModel
+import com.android.customization.picker.common.ui.view.ItemSpacing
+import com.android.wallpaper.R
+import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter
+import com.android.wallpaper.picker.option.ui.binder.OptionItemBinder
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.launch
+
+object GridScreenBinder {
+    fun bind(
+        view: View,
+        viewModel: GridScreenViewModel,
+        lifecycleOwner: LifecycleOwner,
+        backgroundDispatcher: CoroutineDispatcher,
+        onOptionsChanged: () -> Unit,
+    ) {
+        val optionView: RecyclerView = view.requireViewById(R.id.options)
+        optionView.layoutManager =
+            LinearLayoutManager(
+                view.context,
+                RecyclerView.HORIZONTAL,
+                /* reverseLayout= */ false,
+            )
+        optionView.addItemDecoration(ItemSpacing(ItemSpacing.ITEM_SPACING_DP))
+        val adapter =
+            OptionItemAdapter(
+                layoutResourceId = R.layout.grid_option_2,
+                lifecycleOwner = lifecycleOwner,
+                backgroundDispatcher = backgroundDispatcher,
+                foregroundTintSpec =
+                    OptionItemBinder.TintSpec(
+                        selectedColor = view.context.getColor(R.color.text_color_primary),
+                        unselectedColor = view.context.getColor(R.color.text_color_secondary),
+                    ),
+                bindIcon = { foregroundView: View, gridIcon: GridIconViewModel ->
+                    val imageView = foregroundView as? ImageView
+                    imageView?.let { GridIconViewBinder.bind(imageView, gridIcon) }
+                }
+            )
+        optionView.adapter = adapter
+
+        lifecycleOwner.lifecycleScope.launch {
+            lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+                launch {
+                    viewModel.optionItems.collect { options ->
+                        adapter.setItems(options)
+                        onOptionsChanged()
+                    }
+                }
+            }
+        }
+    }
+}
diff --git a/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt b/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt
new file mode 100644
index 00000000..d8cad828
--- /dev/null
+++ b/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt
@@ -0,0 +1,132 @@
+/*
+ * 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.grid.ui.fragment
+
+import android.os.Bundle
+import android.view.LayoutInflater
+import android.view.View
+import android.view.ViewGroup
+import androidx.lifecycle.ViewModelProvider
+import com.android.customization.model.grid.ui.binder.GridScreenBinder
+import com.android.customization.model.grid.ui.viewmodel.GridScreenViewModel
+import com.android.customization.module.ThemePickerInjector
+import com.android.wallpaper.R
+import com.android.wallpaper.module.CurrentWallpaperInfoFactory
+import com.android.wallpaper.module.CustomizationSections
+import com.android.wallpaper.module.InjectorProvider
+import com.android.wallpaper.picker.AppbarFragment
+import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
+import com.android.wallpaper.picker.customization.ui.binder.ScreenPreviewBinder
+import com.android.wallpaper.picker.customization.ui.viewmodel.ScreenPreviewViewModel
+import com.android.wallpaper.util.PreviewUtils
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.suspendCancellableCoroutine
+
+@OptIn(ExperimentalCoroutinesApi::class)
+class GridFragment2 : AppbarFragment() {
+
+    override fun onCreateView(
+        inflater: LayoutInflater,
+        container: ViewGroup?,
+        savedInstanceState: Bundle?
+    ): View {
+        val view =
+            inflater.inflate(
+                R.layout.fragment_grid,
+                container,
+                false,
+            )
+        setUpToolbar(view)
+
+        val injector = InjectorProvider.getInjector() as ThemePickerInjector
+
+        val wallpaperInfoFactory = injector.getCurrentWallpaperInfoFactory(requireContext())
+        var screenPreviewBinding =
+            bindScreenPreview(
+                view,
+                wallpaperInfoFactory,
+                injector.getWallpaperInteractor(requireContext())
+            )
+
+        val viewModelFactory = injector.getGridScreenViewModelFactory(requireContext())
+        GridScreenBinder.bind(
+            view = view,
+            viewModel =
+                ViewModelProvider(
+                    this,
+                    viewModelFactory,
+                )[GridScreenViewModel::class.java],
+            lifecycleOwner = this,
+            backgroundDispatcher = Dispatchers.IO,
+            onOptionsChanged = {
+                screenPreviewBinding.destroy()
+                screenPreviewBinding =
+                    bindScreenPreview(
+                        view,
+                        wallpaperInfoFactory,
+                        injector.getWallpaperInteractor(requireContext())
+                    )
+            }
+        )
+
+        return view
+    }
+
+    override fun getDefaultTitle(): CharSequence {
+        return getString(R.string.grid_title)
+    }
+
+    private fun bindScreenPreview(
+        view: View,
+        wallpaperInfoFactory: CurrentWallpaperInfoFactory,
+        wallpaperInteractor: WallpaperInteractor,
+    ): ScreenPreviewBinder.Binding {
+        return ScreenPreviewBinder.bind(
+            activity = requireActivity(),
+            previewView = view.requireViewById(R.id.preview),
+            viewModel =
+                ScreenPreviewViewModel(
+                    previewUtils =
+                        PreviewUtils(
+                            context = requireContext(),
+                            authorityMetadataKey =
+                                requireContext()
+                                    .getString(
+                                        R.string.grid_control_metadata_name,
+                                    ),
+                        ),
+                    wallpaperInfoProvider = {
+                        suspendCancellableCoroutine { continuation ->
+                            wallpaperInfoFactory.createCurrentWallpaperInfos(
+                                { homeWallpaper, lockWallpaper, _ ->
+                                    continuation.resume(homeWallpaper ?: lockWallpaper, null)
+                                },
+                                /* forceRefresh= */ true,
+                            )
+                        }
+                    },
+                    wallpaperInteractor = wallpaperInteractor,
+                ),
+            lifecycleOwner = this,
+            offsetToStart = false,
+            screen = CustomizationSections.Screen.HOME_SCREEN,
+            onPreviewDirty = { activity?.recreate() },
+        )
+    }
+}
diff --git a/src/com/android/customization/model/grid/ui/viewmodel/GridIconViewModel.kt b/src/com/android/customization/model/grid/ui/viewmodel/GridIconViewModel.kt
new file mode 100644
index 00000000..3942d7cc
--- /dev/null
+++ b/src/com/android/customization/model/grid/ui/viewmodel/GridIconViewModel.kt
@@ -0,0 +1,24 @@
+/*
+ * 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.grid.ui.viewmodel
+
+data class GridIconViewModel(
+    val columns: Int,
+    val rows: Int,
+    val path: String,
+)
diff --git a/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModel.kt b/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModel.kt
new file mode 100644
index 00000000..c11a5947
--- /dev/null
+++ b/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModel.kt
@@ -0,0 +1,106 @@
+/*
+ * 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.grid.ui.viewmodel
+
+import android.annotation.SuppressLint
+import android.content.Context
+import android.content.res.Resources
+import androidx.lifecycle.ViewModel
+import androidx.lifecycle.ViewModelProvider
+import androidx.lifecycle.viewModelScope
+import com.android.customization.model.ResourceConstants
+import com.android.customization.model.grid.domain.interactor.GridInteractor
+import com.android.customization.model.grid.shared.model.GridOptionItemsModel
+import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.launch
+
+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
+
+    val optionItems: Flow<List<OptionItemViewModel<GridIconViewModel>>> =
+        interactor.options.map { model -> toViewModel(model) }
+
+    private fun toViewModel(
+        model: GridOptionItemsModel,
+    ): List<OptionItemViewModel<GridIconViewModel>> {
+        val iconShapePath =
+            applicationContext.resources.getString(
+                Resources.getSystem()
+                    .getIdentifier(
+                        ResourceConstants.CONFIG_ICON_MASK,
+                        "string",
+                        ResourceConstants.ANDROID_PACKAGE,
+                    )
+            )
+
+        return when (model) {
+            is GridOptionItemsModel.Loaded ->
+                model.options.map { option ->
+                    val text = Text.Loaded(option.name)
+                    OptionItemViewModel<GridIconViewModel>(
+                        key =
+                            MutableStateFlow("${option.cols}x${option.rows}") as StateFlow<String>,
+                        payload =
+                            GridIconViewModel(
+                                columns = option.cols,
+                                rows = option.rows,
+                                path = iconShapePath,
+                            ),
+                        text = text,
+                        isSelected = option.isSelected,
+                        onClicked =
+                            option.isSelected.map { isSelected ->
+                                if (!isSelected) {
+                                    { viewModelScope.launch { option.onSelected() } }
+                                } else {
+                                    null
+                                }
+                            },
+                    )
+                }
+            is GridOptionItemsModel.Error -> emptyList()
+        }
+    }
+
+    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
+        }
+    }
+}
diff --git a/src/com/android/customization/model/mode/DarkModeSectionController.java b/src/com/android/customization/model/mode/DarkModeSectionController.java
index f56b7092..ebeaa567 100644
--- a/src/com/android/customization/model/mode/DarkModeSectionController.java
+++ b/src/com/android/customization/model/mode/DarkModeSectionController.java
@@ -59,12 +59,17 @@ public class DarkModeSectionController implements
 
     private Context mContext;
     private DarkModeSectionView mDarkModeSectionView;
+    private final DarkModeSnapshotRestorer mSnapshotRestorer;
 
-    public DarkModeSectionController(Context context, Lifecycle lifecycle) {
+    public DarkModeSectionController(
+            Context context,
+            Lifecycle lifecycle,
+            DarkModeSnapshotRestorer snapshotRestorer) {
         mContext = context;
         mLifecycle = lifecycle;
         mPowerManager = context.getSystemService(PowerManager.class);
         mLifecycle.addObserver(this);
+        mSnapshotRestorer = snapshotRestorer;
     }
 
     @OnLifecycleEvent(Lifecycle.Event.ON_START)
@@ -132,6 +137,7 @@ public class DarkModeSectionController implements
                     mDarkModeSectionView.announceForAccessibility(
                             context.getString(R.string.mode_changed));
                     uiModeManager.setNightModeActivated(viewActivated);
+                    mSnapshotRestorer.store(viewActivated);
                 },
                 /* delayMillis= */ shortDelay);
     }
diff --git a/src/com/android/customization/model/mode/DarkModeSnapshotRestorer.kt b/src/com/android/customization/model/mode/DarkModeSnapshotRestorer.kt
new file mode 100644
index 00000000..93bd0bfd
--- /dev/null
+++ b/src/com/android/customization/model/mode/DarkModeSnapshotRestorer.kt
@@ -0,0 +1,101 @@
+/*
+ * 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.mode
+
+import android.app.UiModeManager
+import android.content.Context
+import android.content.res.Configuration
+import androidx.annotation.VisibleForTesting
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotStore
+import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.withContext
+
+class DarkModeSnapshotRestorer : SnapshotRestorer {
+
+    private val backgroundDispatcher: CoroutineDispatcher
+    private val isActive: () -> Boolean
+    private val setActive: suspend (Boolean) -> Unit
+
+    private var store: SnapshotStore = SnapshotStore.NOOP
+
+    constructor(
+        context: Context,
+        manager: UiModeManager,
+        backgroundDispatcher: CoroutineDispatcher,
+    ) : this(
+        backgroundDispatcher = backgroundDispatcher,
+        isActive = {
+            context.applicationContext.resources.configuration.uiMode and
+                Configuration.UI_MODE_NIGHT_YES != 0
+        },
+        setActive = { isActive -> manager.setNightModeActivated(isActive) },
+    )
+
+    @VisibleForTesting
+    constructor(
+        backgroundDispatcher: CoroutineDispatcher,
+        isActive: () -> Boolean,
+        setActive: suspend (Boolean) -> Unit,
+    ) {
+        this.backgroundDispatcher = backgroundDispatcher
+        this.isActive = isActive
+        this.setActive = setActive
+    }
+
+    override suspend fun setUpSnapshotRestorer(store: SnapshotStore): RestorableSnapshot {
+        this.store = store
+        return snapshot(
+            isActivated = isActive(),
+        )
+    }
+
+    override suspend fun restoreToSnapshot(snapshot: RestorableSnapshot) {
+        val isActivated = snapshot.args[KEY]?.toBoolean() == true
+        withContext(backgroundDispatcher) { setActive(isActivated) }
+    }
+
+    fun store(
+        isActivated: Boolean,
+    ) {
+        store.store(
+            snapshot(
+                isActivated = isActivated,
+            ),
+        )
+    }
+
+    private fun snapshot(
+        isActivated: Boolean,
+    ): RestorableSnapshot {
+        return RestorableSnapshot(
+            args =
+                buildMap {
+                    put(
+                        KEY,
+                        isActivated.toString(),
+                    )
+                }
+        )
+    }
+
+    companion object {
+        private const val KEY = "is_activated"
+    }
+}
diff --git a/src/com/android/customization/model/themedicon/ThemedIconSectionController.java b/src/com/android/customization/model/themedicon/ThemedIconSectionController.java
index a1623d18..5d551a6a 100644
--- a/src/com/android/customization/model/themedicon/ThemedIconSectionController.java
+++ b/src/com/android/customization/model/themedicon/ThemedIconSectionController.java
@@ -20,11 +20,13 @@ import android.os.Bundle;
 import android.view.LayoutInflater;
 
 import androidx.annotation.Nullable;
+import androidx.lifecycle.Observer;
 
+import com.android.customization.model.themedicon.domain.interactor.ThemedIconInteractor;
+import com.android.customization.model.themedicon.domain.interactor.ThemedIconSnapshotRestorer;
 import com.android.customization.picker.themedicon.ThemedIconSectionView;
 import com.android.wallpaper.R;
 import com.android.wallpaper.model.CustomizationSectionController;
-import com.android.wallpaper.model.WorkspaceViewModel;
 
 /** The {@link CustomizationSectionController} for themed icon section. */
 public class ThemedIconSectionController implements
@@ -33,16 +35,26 @@ public class ThemedIconSectionController implements
     private static final String KEY_THEMED_ICON_ENABLED = "SAVED_THEMED_ICON_ENABLED";
 
     private final ThemedIconSwitchProvider mThemedIconOptionsProvider;
-    private final WorkspaceViewModel mWorkspaceViewModel;
+    private final ThemedIconInteractor mInteractor;
+    private final ThemedIconSnapshotRestorer mSnapshotRestorer;
+    private final Observer<Boolean> mIsActivatedChangeObserver;
 
     private ThemedIconSectionView mThemedIconSectionView;
     private boolean mSavedThemedIconEnabled = false;
 
-
-    public ThemedIconSectionController(ThemedIconSwitchProvider themedIconOptionsProvider,
-            WorkspaceViewModel workspaceViewModel, @Nullable Bundle savedInstanceState) {
+    public ThemedIconSectionController(
+            ThemedIconSwitchProvider themedIconOptionsProvider,
+            ThemedIconInteractor interactor,
+            @Nullable Bundle savedInstanceState,
+            ThemedIconSnapshotRestorer snapshotRestorer) {
         mThemedIconOptionsProvider = themedIconOptionsProvider;
-        mWorkspaceViewModel = workspaceViewModel;
+        mInteractor = interactor;
+        mSnapshotRestorer = snapshotRestorer;
+        mIsActivatedChangeObserver = isActivated -> {
+            if (mThemedIconSectionView.isAttachedToWindow()) {
+                mThemedIconSectionView.getSwitch().setChecked(isActivated);
+            }
+        };
 
         if (savedInstanceState != null) {
             mSavedThemedIconEnabled = savedInstanceState.getBoolean(
@@ -64,15 +76,22 @@ public class ThemedIconSectionController implements
         mThemedIconSectionView.getSwitch().setChecked(mSavedThemedIconEnabled);
         mThemedIconOptionsProvider.fetchThemedIconEnabled(
                 enabled -> mThemedIconSectionView.getSwitch().setChecked(enabled));
+        mInteractor.isActivatedAsLiveData().observeForever(mIsActivatedChangeObserver);
         return mThemedIconSectionView;
     }
 
+    @Override
+    public void release() {
+        mInteractor.isActivatedAsLiveData().removeObserver(mIsActivatedChangeObserver);
+    }
+
     private void onViewActivated(Context context, boolean viewActivated) {
         if (context == null) {
             return;
         }
         mThemedIconOptionsProvider.setThemedIconEnabled(viewActivated);
-        mWorkspaceViewModel.getUpdateWorkspace().setValue(viewActivated);
+        mInteractor.setActivated(viewActivated);
+        mSnapshotRestorer.store(viewActivated);
     }
 
     @Override
diff --git a/src/com/android/customization/model/themedicon/ThemedIconSwitchProvider.java b/src/com/android/customization/model/themedicon/ThemedIconSwitchProvider.java
index 9acd3190..5e2a60a1 100644
--- a/src/com/android/customization/model/themedicon/ThemedIconSwitchProvider.java
+++ b/src/com/android/customization/model/themedicon/ThemedIconSwitchProvider.java
@@ -118,7 +118,7 @@ public class ThemedIconSwitchProvider {
      *
      * <p>The value would also be stored in SharedPreferences.
      */
-    protected void setThemedIconEnabled(boolean enabled) {
+    public void setThemedIconEnabled(boolean enabled) {
         mExecutorService.submit(() -> {
             ContentValues values = new ContentValues();
             values.put(COL_ICON_THEMED_VALUE, enabled);
diff --git a/src/com/android/customization/model/themedicon/data/repository/ThemedIconRepository.kt b/src/com/android/customization/model/themedicon/data/repository/ThemedIconRepository.kt
new file mode 100644
index 00000000..91088111
--- /dev/null
+++ b/src/com/android/customization/model/themedicon/data/repository/ThemedIconRepository.kt
@@ -0,0 +1,30 @@
+/*
+ * 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.themedicon.data.repository
+
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.asStateFlow
+
+class ThemeIconRepository {
+    private val _isActivated = MutableStateFlow(false)
+    val isActivated = _isActivated.asStateFlow()
+
+    fun setActivated(isActivated: Boolean) {
+        _isActivated.value = isActivated
+    }
+}
diff --git a/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractor.kt b/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractor.kt
new file mode 100644
index 00000000..1cfe8776
--- /dev/null
+++ b/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractor.kt
@@ -0,0 +1,38 @@
+/*
+ * 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.themedicon.domain.interactor
+
+import androidx.lifecycle.LiveData
+import androidx.lifecycle.asLiveData
+import com.android.customization.model.themedicon.data.repository.ThemeIconRepository
+
+class ThemedIconInteractor(
+    private val repository: ThemeIconRepository,
+) {
+    val isActivated = repository.isActivated
+
+    private var isActivatedAsLiveData: LiveData<Boolean>? = null
+
+    fun isActivatedAsLiveData(): LiveData<Boolean> {
+        return isActivatedAsLiveData ?: isActivated.asLiveData().also { isActivatedAsLiveData = it }
+    }
+
+    fun setActivated(isActivated: Boolean) {
+        repository.setActivated(isActivated)
+    }
+}
diff --git a/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorer.kt b/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorer.kt
new file mode 100644
index 00000000..cacc45ef
--- /dev/null
+++ b/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorer.kt
@@ -0,0 +1,60 @@
+/*
+ * 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.themedicon.domain.interactor
+
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotStore
+import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
+
+class ThemedIconSnapshotRestorer(
+    private val isActivated: () -> Boolean,
+    private val setActivated: (isActivated: Boolean) -> Unit,
+    private val interactor: ThemedIconInteractor,
+) : SnapshotRestorer {
+
+    private var store: SnapshotStore = SnapshotStore.NOOP
+
+    override suspend fun setUpSnapshotRestorer(store: SnapshotStore): RestorableSnapshot {
+        this.store = store
+        return snapshot()
+    }
+
+    override suspend fun restoreToSnapshot(snapshot: RestorableSnapshot) {
+        val isActivated = snapshot.args[KEY]?.toBoolean() == true
+        setActivated(isActivated)
+        interactor.setActivated(isActivated)
+    }
+
+    fun store(
+        isActivated: Boolean,
+    ) {
+        store.store(snapshot(isActivated = isActivated))
+    }
+
+    private fun snapshot(
+        isActivated: Boolean? = null,
+    ): RestorableSnapshot {
+        return RestorableSnapshot(
+            args = buildMap { put(KEY, (isActivated ?: isActivated()).toString()) }
+        )
+    }
+
+    companion object {
+        private const val KEY = "is_activated"
+    }
+}