Merge "Fix preview crash on unfold in multi-crop" into main
diff --git a/src/com/android/wallpaper/picker/CustomizationPickerActivity.java b/src/com/android/wallpaper/picker/CustomizationPickerActivity.java
index cf705e0..9bdc85b 100644
--- a/src/com/android/wallpaper/picker/CustomizationPickerActivity.java
+++ b/src/com/android/wallpaper/picker/CustomizationPickerActivity.java
@@ -92,7 +92,7 @@
         mNetworkStatus = mNetworkStatusNotifier.getNetworkStatus();
         mDisplayUtils = injector.getDisplayUtils(this);
 
-        enforceOrientation();
+        enforcePortraitForHandheldAndFoldedDisplay();
 
         // Restore this Activity's state before restoring contained Fragments state.
         super.onCreate(savedInstanceState);
@@ -394,22 +394,21 @@
     @Override
     public void onConfigurationChanged(@NonNull Configuration newConfig) {
         super.onConfigurationChanged(newConfig);
-        enforceOrientation();
+        enforcePortraitForHandheldAndFoldedDisplay();
     }
 
     /**
-     * Allows any orientation for large screen devices (tablets and unfolded foldables) while
-     * forcing portrait for smaller screens (handheld and folded foldables).
+     * If the display is a handheld display or a folded display from a foldable, we enforce the
+     * activity to be portrait.
      *
      * This method should be called upon initialization of this activity, and whenever there is a
      * configuration change.
      */
     @SuppressLint("SourceLockedOrientationActivity")
-    private void enforceOrientation() {
-        int wantedOrientation =
-                mDisplayUtils.isLargeScreenDevice() && mDisplayUtils.isOnWallpaperDisplay(this)
-                        ? ActivityInfo.SCREEN_ORIENTATION_UNSPECIFIED
-                        : ActivityInfo.SCREEN_ORIENTATION_PORTRAIT;
+    private void enforcePortraitForHandheldAndFoldedDisplay() {
+        int wantedOrientation = mDisplayUtils.isLargeScreenOrUnfoldedDisplay(this)
+                ? ActivityInfo.SCREEN_ORIENTATION_UNSPECIFIED
+                : ActivityInfo.SCREEN_ORIENTATION_PORTRAIT;
         if (getRequestedOrientation() != wantedOrientation) {
             setRequestedOrientation(wantedOrientation);
         }
diff --git a/src/com/android/wallpaper/picker/customization/data/content/WallpaperClient.kt b/src/com/android/wallpaper/picker/customization/data/content/WallpaperClient.kt
index 7ca9791..649a140 100644
--- a/src/com/android/wallpaper/picker/customization/data/content/WallpaperClient.kt
+++ b/src/com/android/wallpaper/picker/customization/data/content/WallpaperClient.kt
@@ -17,6 +17,7 @@
 
 package com.android.wallpaper.picker.customization.data.content
 
+import android.app.WallpaperColors
 import android.app.WallpaperManager
 import android.graphics.Bitmap
 import android.graphics.Point
@@ -99,4 +100,7 @@
         displaySizes: List<Point>,
         @WallpaperManager.SetWallpaperFlags which: Int
     ): Map<Point, Rect>?
+
+    /** Returns the wallpaper colors for preview a bitmap with a set of crop hints */
+    suspend fun getWallpaperColors(bitmap: Bitmap, cropHints: Map<Point, Rect>?): WallpaperColors?
 }
diff --git a/src/com/android/wallpaper/picker/customization/data/content/WallpaperClientImpl.kt b/src/com/android/wallpaper/picker/customization/data/content/WallpaperClientImpl.kt
index b40b083..e9f97b0 100644
--- a/src/com/android/wallpaper/picker/customization/data/content/WallpaperClientImpl.kt
+++ b/src/com/android/wallpaper/picker/customization/data/content/WallpaperClientImpl.kt
@@ -17,6 +17,7 @@
 
 package com.android.wallpaper.picker.customization.data.content
 
+import android.app.WallpaperColors
 import android.app.WallpaperManager
 import android.app.WallpaperManager.FLAG_LOCK
 import android.app.WallpaperManager.FLAG_SYSTEM
@@ -45,6 +46,7 @@
 import com.android.wallpaper.module.logging.UserEventLogger.SetWallpaperEntryPoint
 import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination
 import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination.BOTH
+import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination.Companion.toDestinationInt
 import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination.HOME
 import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination.LOCK
 import com.android.wallpaper.picker.customization.shared.model.WallpaperModel
@@ -374,7 +376,7 @@
         val uri =
             Uri.parse(uriString)
                 ?.buildUpon()
-                ?.appendQueryParameter("destination", destination.toString())
+                ?.appendQueryParameter("destination", destination.toDestinationInt().toString())
                 ?.build()
                 ?: return null
         val authority = uri.authority ?: return null
@@ -570,6 +572,13 @@
         return cropHintsMap
     }
 
+    override suspend fun getWallpaperColors(
+        bitmap: Bitmap,
+        cropHints: Map<Point, Rect>?
+    ): WallpaperColors? {
+        return wallpaperManager.getWallpaperColors(bitmap, cropHints)
+    }
+
     fun WallpaperDestination.asString(): String {
         return when (this) {
             BOTH -> SCREEN_ALL
@@ -593,21 +602,26 @@
      * on the view size hosting the preview and the wallpaper zoom of the preview on that view,
      * whereas the rest of multi-crop is based on full wallpaper size. So scaled back at the end.
      *
+     * If [CropSizeModel] is null, returns the original cropHint without parallax.
+     *
      * @param wallpaperSize full wallpaper image size.
      */
     private fun FullPreviewCropModel.adjustCropForParallax(
         wallpaperSize: Point,
     ): Rect {
-        return WallpaperCropUtils.calculateCropRect(
-                context,
-                hostViewSize,
-                cropSurfaceSize,
-                wallpaperSize,
-                cropHint,
-                wallpaperZoom,
-                /* cropExtraWidth= */ true,
-            )
-            .apply { scale(1f / wallpaperZoom) }
+        return cropSizeModel?.let {
+            WallpaperCropUtils.calculateCropRect(
+                    context,
+                    it.hostViewSize,
+                    it.cropSurfaceSize,
+                    wallpaperSize,
+                    cropHint,
+                    it.wallpaperZoom,
+                    /* cropExtraWidth= */ true,
+                )
+                .apply { scale(1f / it.wallpaperZoom) }
+        }
+            ?: cropHint
     }
 
     companion object {
diff --git a/src/com/android/wallpaper/picker/customization/data/repository/WallpaperRepository.kt b/src/com/android/wallpaper/picker/customization/data/repository/WallpaperRepository.kt
index 669f5d0..ec3e58c 100644
--- a/src/com/android/wallpaper/picker/customization/data/repository/WallpaperRepository.kt
+++ b/src/com/android/wallpaper/picker/customization/data/repository/WallpaperRepository.kt
@@ -17,8 +17,10 @@
 
 package com.android.wallpaper.picker.customization.data.repository
 
+import android.app.WallpaperColors
 import android.graphics.Bitmap
 import android.graphics.Point
+import android.graphics.Rect
 import android.util.LruCache
 import com.android.wallpaper.module.WallpaperPreferences
 import com.android.wallpaper.module.logging.UserEventLogger.SetWallpaperEntryPoint
@@ -175,6 +177,9 @@
         }
     }
 
+    suspend fun getWallpaperColors(bitmap: Bitmap, cropHints: Map<Point, Rect>?): WallpaperColors? =
+        withContext(backgroundDispatcher) { client.getWallpaperColors(bitmap, cropHints) }
+
     companion object {
         const val DEFAULT_KEY = "default_missing_key"
         /** The maximum number of options to show, including the currently-selected one. */
diff --git a/src/com/android/wallpaper/picker/customization/shared/model/WallpaperDestination.kt b/src/com/android/wallpaper/picker/customization/shared/model/WallpaperDestination.kt
index 6adbb22..4b8c680 100644
--- a/src/com/android/wallpaper/picker/customization/shared/model/WallpaperDestination.kt
+++ b/src/com/android/wallpaper/picker/customization/shared/model/WallpaperDestination.kt
@@ -20,6 +20,10 @@
 import android.app.WallpaperManager.FLAG_LOCK
 import android.app.WallpaperManager.FLAG_SYSTEM
 import android.app.WallpaperManager.SetWallpaperFlags
+import com.android.wallpaper.module.WallpaperPersister.DEST_BOTH
+import com.android.wallpaper.module.WallpaperPersister.DEST_HOME_SCREEN
+import com.android.wallpaper.module.WallpaperPersister.DEST_LOCK_SCREEN
+import com.android.wallpaper.module.WallpaperPersister.Destination
 
 /** Enumerates all known wallpaper destinations. */
 enum class WallpaperDestination {
@@ -39,5 +43,14 @@
                 else -> throw IllegalArgumentException("Bad @SetWallpaperFlags value $flags")
             }
         }
+
+        @Destination
+        fun WallpaperDestination.toDestinationInt(): Int {
+            return when (this) {
+                BOTH -> DEST_BOTH
+                HOME -> DEST_HOME_SCREEN
+                LOCK -> DEST_LOCK_SCREEN
+            }
+        }
     }
 }
diff --git a/src/com/android/wallpaper/picker/preview/domain/interactor/WallpaperPreviewInteractor.kt b/src/com/android/wallpaper/picker/preview/domain/interactor/WallpaperPreviewInteractor.kt
index f243cec..20da4f4 100644
--- a/src/com/android/wallpaper/picker/preview/domain/interactor/WallpaperPreviewInteractor.kt
+++ b/src/com/android/wallpaper/picker/preview/domain/interactor/WallpaperPreviewInteractor.kt
@@ -16,8 +16,10 @@
 
 package com.android.wallpaper.picker.preview.domain.interactor
 
+import android.app.WallpaperColors
 import android.graphics.Bitmap
 import android.graphics.Point
+import android.graphics.Rect
 import com.android.wallpaper.module.logging.UserEventLogger
 import com.android.wallpaper.picker.customization.data.repository.WallpaperRepository
 import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination
@@ -73,4 +75,7 @@
             wallpaperModel,
         )
     }
+
+    suspend fun getWallpaperColors(bitmap: Bitmap, cropHints: Map<Point, Rect>?): WallpaperColors? =
+        wallpaperRepository.getWallpaperColors(bitmap, cropHints)
 }
diff --git a/src/com/android/wallpaper/picker/preview/shared/model/FullPreviewCropModel.kt b/src/com/android/wallpaper/picker/preview/shared/model/FullPreviewCropModel.kt
index dbc2751..6e83b21 100644
--- a/src/com/android/wallpaper/picker/preview/shared/model/FullPreviewCropModel.kt
+++ b/src/com/android/wallpaper/picker/preview/shared/model/FullPreviewCropModel.kt
@@ -19,14 +19,30 @@
 import android.graphics.Point
 import android.graphics.Rect
 
-/** Data class for cropHints related info. */
+/**
+ * Data class represents user's cropHint for a dimension.
+ *
+ * It could be one of below:
+ * 1. A current wallpaper crop.
+ * 2. User's crop via full preview.
+ * 3. Default crop from small preview.
+ *
+ * Only #2 will it contains [cropSizeModel], the other cases parallax (0 for #3) has already
+ * included in [cropHint].
+ */
 data class FullPreviewCropModel(
-    /** The user's crop of wallpaper on FullPreviewFragment wrt the full wallpaper size. */
-    val cropHint: Rect = Rect(0, 0, 0, 0),
+    /** The user's crop of wallpaper based on the full wallpaper size. */
+    val cropHint: Rect,
+    /** The data required to compute parallax for this crop, null for no parallax. */
+    val cropSizeModel: CropSizeModel?,
+)
+
+/** Required for computing parallax. */
+data class CropSizeModel(
     /** The zoom of the wallpaper on its hosting view when user selects the cropHint. */
-    val wallpaperZoom: Float = 0f,
+    val wallpaperZoom: Float,
     /** The size of the view hosting the wallpaper, e.g. SurfaceView. */
-    val hostViewSize: Point = Point(0, 0),
+    val hostViewSize: Point,
     /** A larger version of hostViewSize that can safely contain parallax. */
-    val cropSurfaceSize: Point = Point(0, 0),
+    val cropSurfaceSize: Point,
 )
diff --git a/src/com/android/wallpaper/picker/preview/ui/WallpaperPreviewActivity.kt b/src/com/android/wallpaper/picker/preview/ui/WallpaperPreviewActivity.kt
index 0560ba5..74e8eff 100644
--- a/src/com/android/wallpaper/picker/preview/ui/WallpaperPreviewActivity.kt
+++ b/src/com/android/wallpaper/picker/preview/ui/WallpaperPreviewActivity.kt
@@ -18,6 +18,7 @@
 import android.content.Context
 import android.content.Intent
 import android.content.pm.ActivityInfo
+import android.content.res.Configuration
 import android.graphics.Color
 import android.os.Bundle
 import android.widget.Toast
@@ -64,6 +65,7 @@
 
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
+        enforcePortraitForHandheldAndFoldedDisplay()
         window.navigationBarColor = Color.TRANSPARENT
         window.statusBarColor = Color.TRANSPARENT
         setContentView(R.layout.activity_wallpaper_preview)
@@ -139,9 +141,6 @@
 
     override fun onResume() {
         super.onResume()
-        requestedOrientation =
-            if (displayUtils.isOnWallpaperDisplay(this)) ActivityInfo.SCREEN_ORIENTATION_USER
-            else ActivityInfo.SCREEN_ORIENTATION_PORTRAIT
         if (isInMultiWindowMode) {
             Toast.makeText(this, R.string.wallpaper_exit_split_screen, Toast.LENGTH_SHORT).show()
             onBackPressedDispatcher.onBackPressed()
@@ -157,6 +156,11 @@
         super.onDestroy()
     }
 
+    override fun onConfigurationChanged(newConfig: Configuration) {
+        super.onConfigurationChanged(newConfig)
+        enforcePortraitForHandheldAndFoldedDisplay()
+    }
+
     private fun WallpaperInfo.convertToWallpaperModel(): WallpaperModel {
         return wallpaperModelFactory.getWallpaperModel(appContext, this)
     }
@@ -189,4 +193,21 @@
             return intent
         }
     }
+
+    /**
+     * If the display is a handheld display or a folded display from a foldable, we enforce the
+     * activity to be portrait.
+     *
+     * This method should be called upon initialization of this activity, and whenever there is a
+     * configuration change.
+     */
+    private fun enforcePortraitForHandheldAndFoldedDisplay() {
+        val wantedOrientation =
+            if (displayUtils.isLargeScreenOrUnfoldedDisplay(this))
+                ActivityInfo.SCREEN_ORIENTATION_UNSPECIFIED
+            else ActivityInfo.SCREEN_ORIENTATION_PORTRAIT
+        if (requestedOrientation != wantedOrientation) {
+            requestedOrientation = wantedOrientation
+        }
+    }
 }
diff --git a/src/com/android/wallpaper/picker/preview/ui/binder/FullWallpaperPreviewBinder.kt b/src/com/android/wallpaper/picker/preview/ui/binder/FullWallpaperPreviewBinder.kt
index e8c458a..f9179cd 100644
--- a/src/com/android/wallpaper/picker/preview/ui/binder/FullWallpaperPreviewBinder.kt
+++ b/src/com/android/wallpaper/picker/preview/ui/binder/FullWallpaperPreviewBinder.kt
@@ -16,7 +16,6 @@
 package com.android.wallpaper.picker.preview.ui.binder
 
 import android.content.Context
-import android.graphics.PointF
 import android.graphics.Rect
 import android.view.LayoutInflater
 import android.view.SurfaceHolder
@@ -31,8 +30,9 @@
 import com.android.wallpaper.picker.TouchForwardingLayout
 import com.android.wallpaper.picker.data.WallpaperModel
 import com.android.wallpaper.picker.di.modules.MainDispatcher
+import com.android.wallpaper.picker.preview.shared.model.CropSizeModel
 import com.android.wallpaper.picker.preview.shared.model.FullPreviewCropModel
-import com.android.wallpaper.picker.preview.ui.util.FullResImageViewUtil.getCropRect
+import com.android.wallpaper.picker.preview.ui.util.SubsamplingScaleImageViewUtil.setOnNewCropListener
 import com.android.wallpaper.picker.preview.ui.util.SurfaceViewUtil
 import com.android.wallpaper.picker.preview.ui.util.SurfaceViewUtil.attachView
 import com.android.wallpaper.picker.preview.ui.view.FullPreviewFrameLayout
@@ -41,7 +41,6 @@
 import com.android.wallpaper.util.WallpaperCropUtils
 import com.android.wallpaper.util.wallpaperconnection.WallpaperConnectionUtils
 import com.davemorrissey.labs.subscaleview.SubsamplingScaleImageView
-import com.davemorrissey.labs.subscaleview.SubsamplingScaleImageView.OnStateChangedListener
 import java.lang.Integer.min
 import kotlin.math.max
 import kotlinx.coroutines.CoroutineScope
@@ -108,12 +107,15 @@
                                             surfaceView,
                                         ) { crop, zoom ->
                                             viewModel.staticWallpaperPreviewViewModel
-                                                .fullPreviewCropModel =
+                                                .fullPreviewCropModels[config.displaySize] =
                                                 FullPreviewCropModel(
                                                     cropHint = crop,
-                                                    wallpaperZoom = zoom,
-                                                    hostViewSize = surfaceSize,
-                                                    cropSurfaceSize = cropSurfaceSize,
+                                                    cropSizeModel =
+                                                        CropSizeModel(
+                                                            wallpaperZoom = zoom,
+                                                            hostViewSize = surfaceSize,
+                                                            cropSurfaceSize = cropSurfaceSize,
+                                                        ),
                                                 )
                                         }
 
@@ -132,7 +134,6 @@
                                         viewModel.staticWallpaperPreviewViewModel,
                                         config.displaySize,
                                         lifecycleOwner,
-                                        allowUserCropping,
                                     )
                                 }
                             }
@@ -170,20 +171,4 @@
         setForwardingEnabled(true)
         setTargetView(targetView)
     }
-
-    private fun SubsamplingScaleImageView.setOnNewCropListener(
-        onNewCrop: (crop: Rect, zoom: Float) -> Unit
-    ) {
-        setOnStateChangedListener(
-            object : OnStateChangedListener {
-                override fun onScaleChanged(p0: Float, p1: Int) {
-                    onNewCrop.invoke(getCropRect(), scale)
-                }
-
-                override fun onCenterChanged(p0: PointF?, p1: Int) {
-                    onNewCrop.invoke(getCropRect(), scale)
-                }
-            }
-        )
-    }
 }
diff --git a/src/com/android/wallpaper/picker/preview/ui/binder/StaticWallpaperPreviewBinder.kt b/src/com/android/wallpaper/picker/preview/ui/binder/StaticWallpaperPreviewBinder.kt
index f3d28de..c481ea7 100644
--- a/src/com/android/wallpaper/picker/preview/ui/binder/StaticWallpaperPreviewBinder.kt
+++ b/src/com/android/wallpaper/picker/preview/ui/binder/StaticWallpaperPreviewBinder.kt
@@ -17,7 +17,6 @@
 
 import android.animation.Animator
 import android.animation.AnimatorListenerAdapter
-import android.graphics.Bitmap
 import android.graphics.Point
 import android.graphics.Rect
 import android.graphics.RenderEffect
@@ -46,18 +45,12 @@
     private val ALPHA_OUT: Interpolator = PathInterpolator(0f, 0f, 0.8f, 1f)
     private const val CROSS_FADE_DURATION: Long = 200
 
-    /**
-     * Binds static wallpaper preview.
-     *
-     * @param fullPreviewCropModel null if this is not binding the full preview.
-     */
     fun bind(
         lowResImageView: ImageView,
         fullResImageView: SubsamplingScaleImageView,
         viewModel: StaticWallpaperPreviewViewModel,
         displaySize: Point,
         viewLifecycleOwner: LifecycleOwner,
-        allowUserCropping: Boolean = false,
         shouldCalibrateWithSystemScale: Boolean = false,
     ) {
         lowResImageView.initLowResImageView()
@@ -71,7 +64,7 @@
                     viewModel.subsamplingScaleImageViewModel.collect { imageModel ->
                         val cropHint = imageModel.fullPreviewCropModels?.get(displaySize)?.cropHint
                         fullResImageView.setFullResImage(
-                            imageModel.rawWallpaperBitmap,
+                            ImageSource.cachedBitmap(imageModel.rawWallpaperBitmap),
                             imageModel.rawWallpaperSize,
                             displaySize,
                             cropHint,
@@ -79,24 +72,22 @@
                             shouldCalibrateWithSystemScale,
                         )
 
-                        if (allowUserCropping) {
-                            viewModel.fullPreviewCropModel?.let {
-                                viewModel.fullPreviewCropModel =
-                                    FullPreviewCropModel(
-                                        cropHint = cropHint
-                                                ?: WallpaperCropUtils.calculateVisibleRect(
-                                                    imageModel.rawWallpaperSize,
-                                                    Point(
-                                                        fullResImageView.measuredWidth,
-                                                        fullResImageView.measuredHeight
-                                                    )
-                                                ),
-                                        it.wallpaperZoom,
-                                        it.hostViewSize,
-                                        it.cropSurfaceSize,
-                                    )
-                            }
-                        }
+                        // Fill in the default crop region if the displaySize for this preview is
+                        // missing.
+                        viewModel.fullPreviewCropModels.putIfAbsent(
+                            displaySize,
+                            FullPreviewCropModel(
+                                cropHint =
+                                    WallpaperCropUtils.calculateVisibleRect(
+                                        imageModel.rawWallpaperSize,
+                                        Point(
+                                            fullResImageView.measuredWidth,
+                                            fullResImageView.measuredHeight
+                                        )
+                                    ),
+                                cropSizeModel = null,
+                            )
+                        )
 
                         crossFadeInFullResImageView(lowResImageView, fullResImageView)
                     }
@@ -128,7 +119,7 @@
      *   this system scale to [SubsamplingScaleImageView].
      */
     private fun SubsamplingScaleImageView.setFullResImage(
-        rawWallpaperBitmap: Bitmap,
+        imageSource: ImageSource,
         rawWallpaperSize: Point,
         displaySize: Point,
         cropHint: Rect?,
@@ -136,7 +127,7 @@
         shouldCalibrateWithSystemScale: Boolean = false,
     ) {
         // Set the full res image
-        setImage(ImageSource.bitmap(rawWallpaperBitmap))
+        setImage(imageSource)
         // Calculate the scale and the center point for the full res image
         FullResImageViewUtil.getScaleAndCenter(
                 Point(measuredWidth, measuredHeight),
diff --git a/src/com/android/wallpaper/picker/preview/ui/util/SubsamplingScaleImageViewUtil.kt b/src/com/android/wallpaper/picker/preview/ui/util/SubsamplingScaleImageViewUtil.kt
new file mode 100644
index 0000000..ae7b3c9
--- /dev/null
+++ b/src/com/android/wallpaper/picker/preview/ui/util/SubsamplingScaleImageViewUtil.kt
@@ -0,0 +1,41 @@
+/*
+ * Copyright (C) 2024 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.wallpaper.picker.preview.ui.util
+
+import android.graphics.PointF
+import android.graphics.Rect
+import com.android.wallpaper.picker.preview.ui.util.FullResImageViewUtil.getCropRect
+import com.davemorrissey.labs.subscaleview.SubsamplingScaleImageView
+
+object SubsamplingScaleImageViewUtil {
+
+    fun SubsamplingScaleImageView.setOnNewCropListener(
+        onNewCrop: (crop: Rect, zoom: Float) -> Unit
+    ) {
+        setOnStateChangedListener(
+            object : SubsamplingScaleImageView.OnStateChangedListener {
+                override fun onScaleChanged(p0: Float, p1: Int) {
+                    onNewCrop.invoke(getCropRect(), scale)
+                }
+
+                override fun onCenterChanged(p0: PointF?, p1: Int) {
+                    onNewCrop.invoke(getCropRect(), scale)
+                }
+            }
+        )
+    }
+}
diff --git a/src/com/android/wallpaper/picker/preview/ui/viewmodel/StaticWallpaperPreviewViewModel.kt b/src/com/android/wallpaper/picker/preview/ui/viewmodel/StaticWallpaperPreviewViewModel.kt
index a56518f..6b64bab 100644
--- a/src/com/android/wallpaper/picker/preview/ui/viewmodel/StaticWallpaperPreviewViewModel.kt
+++ b/src/com/android/wallpaper/picker/preview/ui/viewmodel/StaticWallpaperPreviewViewModel.kt
@@ -21,6 +21,7 @@
 import android.graphics.BitmapFactory
 import android.graphics.ColorSpace
 import android.graphics.Point
+import android.graphics.Rect
 import com.android.wallpaper.asset.Asset
 import com.android.wallpaper.asset.StreamableAsset
 import com.android.wallpaper.module.WallpaperPreferences
@@ -37,13 +38,16 @@
 import javax.inject.Inject
 import kotlinx.coroutines.CancellableContinuation
 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.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.suspendCancellableCoroutine
 
 /** View model for static wallpaper preview used in [WallpaperPreviewActivity] and its fragments */
@@ -51,18 +55,35 @@
 class StaticWallpaperPreviewViewModel
 @Inject
 constructor(
-    private val interactor: WallpaperPreviewInteractor,
+    interactor: WallpaperPreviewInteractor,
     @ApplicationContext private val context: Context,
     private val wallpaperPreferences: WallpaperPreferences,
     @BackgroundDispatcher private val bgDispatcher: CoroutineDispatcher,
+    viewModelScope: CoroutineScope,
 ) {
-    /** The state of static wallpaper crop in full preview, before user confirmation. */
-    var fullPreviewCropModel: FullPreviewCropModel? = null
+    /**
+     * The state of static wallpaper crop in full preview, before user confirmation.
+     *
+     * The initial value should be the default crop on small preview, which could be the cropHints
+     * for current wallpaper or default crop area for a new wallpaper.
+     */
+    val fullPreviewCropModels: MutableMap<Point, FullPreviewCropModel> = mutableMapOf()
 
-    /** The info picker needs to post process crops for setting static wallpaper. */
+    /**
+     * The info picker needs to post process crops for setting static wallpaper.
+     *
+     * It will be filled with current cropHints when previewing current wallpaper, and null when
+     * previewing a new wallpaper, and gets updated through [updateCropHintsInfo] when user picks a
+     * new crop.
+     */
     private val cropHintsInfo: MutableStateFlow<Map<Point, FullPreviewCropModel>?> =
         MutableStateFlow(null)
 
+    private val cropHints: Flow<Map<Point, Rect>?> =
+        cropHintsInfo.map { cropHintsInfoMap ->
+            cropHintsInfoMap?.map { entry -> entry.key to entry.value.cropHint }?.toMap()
+        }
+
     val staticWallpaperModel: Flow<StaticWallpaperModel> =
         interactor.wallpaperModel.map { it as? StaticWallpaperModel }.filterNotNull()
     val lowResBitmap: Flow<Bitmap> =
@@ -85,6 +106,10 @@
                 }
             }
             .flowOn(bgDispatcher)
+            // We only want to decode bitmap every time when wallpaper model is updated, instead of
+            // a new subscriber listens to this flow. So we need to use shareIn.
+            .shareIn(viewModelScope, SharingStarted.Lazily, 1)
+
     val fullResWallpaperViewModel: Flow<FullResWallpaperViewModel?> =
         combine(assetDetail, cropHintsInfo) { assetDetail, cropHintsInfo ->
                 if (assetDetail == null) {
@@ -99,15 +124,35 @@
             .flowOn(bgDispatcher)
     val subsamplingScaleImageViewModel: Flow<FullResWallpaperViewModel> =
         fullResWallpaperViewModel.filterNotNull()
-    val wallpaperColors: Flow<WallpaperColorsModel> =
+    // TODO (b/315856338): cache wallpaper colors in preferences
+    private val storedWallpaperColors: Flow<WallpaperColors?> =
         staticWallpaperModel
-            .map {
-                WallpaperColorsModel.Loaded(
-                    wallpaperPreferences.getWallpaperColors(it.commonWallpaperData.id.uniqueId)
-                )
-            }
+            .map { wallpaperPreferences.getWallpaperColors(it.commonWallpaperData.id.uniqueId) }
             .distinctUntilChanged()
+    val wallpaperColors: Flow<WallpaperColorsModel> =
+        combine(storedWallpaperColors, subsamplingScaleImageViewModel, cropHints) {
+            storedColors,
+            wallpaperViewModel,
+            cropHints ->
+            WallpaperColorsModel.Loaded(
+                if (cropHints == null) {
+                    storedColors
+                        ?: interactor.getWallpaperColors(
+                            wallpaperViewModel.rawWallpaperBitmap,
+                            null
+                        )
+                } else {
+                    interactor.getWallpaperColors(wallpaperViewModel.rawWallpaperBitmap, cropHints)
+                }
+            )
+        }
 
+    /**
+     * Updates new cropHints per displaySize that's been confirmed by the user.
+     *
+     * That's when picker gets current cropHints from [WallpaperManager] or when user crops and
+     * confirms a crop.
+     */
     fun updateCropHintsInfo(cropHintsInfo: Map<Point, FullPreviewCropModel>) {
         this.cropHintsInfo.value = this.cropHintsInfo.value?.plus(cropHintsInfo) ?: cropHintsInfo
     }
@@ -124,7 +169,7 @@
     private suspend fun Asset.decodeBitmap(dimensions: Point): Bitmap? =
         suspendCancellableCoroutine { k: CancellableContinuation<Bitmap?> ->
             val callback = Asset.BitmapReceiver { k.resumeWith(Result.success(it)) }
-            decodeBitmap(dimensions.x, dimensions.y, false, callback)
+            decodeBitmap(dimensions.x, dimensions.y, /* hardwareBitmapAllowed= */ false, callback)
         }
 
     private suspend fun Asset.getStream(): InputStream? =
@@ -157,4 +202,23 @@
         }
         return colors
     }
+
+    class Factory
+    @Inject
+    constructor(
+        private val interactor: WallpaperPreviewInteractor,
+        @ApplicationContext private val context: Context,
+        private val wallpaperPreferences: WallpaperPreferences,
+        @BackgroundDispatcher private val bgDispatcher: CoroutineDispatcher,
+    ) {
+        fun create(viewModelScope: CoroutineScope): StaticWallpaperPreviewViewModel {
+            return StaticWallpaperPreviewViewModel(
+                interactor = interactor,
+                context = context,
+                wallpaperPreferences = wallpaperPreferences,
+                bgDispatcher = bgDispatcher,
+                viewModelScope = viewModelScope,
+            )
+        }
+    }
 }
diff --git a/src/com/android/wallpaper/picker/preview/ui/viewmodel/WallpaperPreviewViewModel.kt b/src/com/android/wallpaper/picker/preview/ui/viewmodel/WallpaperPreviewViewModel.kt
index 4326e44..3cd5286 100644
--- a/src/com/android/wallpaper/picker/preview/ui/viewmodel/WallpaperPreviewViewModel.kt
+++ b/src/com/android/wallpaper/picker/preview/ui/viewmodel/WallpaperPreviewViewModel.kt
@@ -57,13 +57,15 @@
 @Inject
 constructor(
     private val interactor: WallpaperPreviewInteractor,
-    val staticWallpaperPreviewViewModel: StaticWallpaperPreviewViewModel,
+    staticWallpaperPreviewViewModelFactory: StaticWallpaperPreviewViewModel.Factory,
     val previewActionsViewModel: PreviewActionsViewModel,
     private val displayUtils: DisplayUtils,
     @HomeScreenPreviewUtils private val homePreviewUtils: PreviewUtils,
     @LockScreenPreviewUtils private val lockPreviewUtils: PreviewUtils,
 ) : ViewModel() {
 
+    val staticWallpaperPreviewViewModel =
+        staticWallpaperPreviewViewModelFactory.create(viewModelScope)
     val smallerDisplaySize = displayUtils.getRealSize(displayUtils.getSmallerDisplay())
     val wallpaperDisplaySize = displayUtils.getRealSize(displayUtils.getWallpaperDisplay())
     var isViewAsHome = false
@@ -101,6 +103,7 @@
                     cropHints.mapValues {
                         FullPreviewCropModel(
                             cropHint = it.value,
+                            cropSizeModel = null,
                         )
                     }
                 )
@@ -163,23 +166,11 @@
         _fullWorkspacePreviewConfigViewModel.filterNotNull()
 
     val onCropButtonClick: Flow<(() -> Unit)?> =
-        combine(wallpaper, fullWallpaperPreviewConfigViewModel.filterNotNull()) {
-            wallpaper,
-            previewViewModel ->
+        combine(wallpaper, fullWallpaperPreviewConfigViewModel.filterNotNull()) { wallpaper, _ ->
             if (wallpaper is StaticWallpaperModel && !wallpaper.isDownloadableWallpaper()) {
                 {
-                    staticWallpaperPreviewViewModel.fullPreviewCropModel?.let {
-                        staticWallpaperPreviewViewModel.updateCropHintsInfo(
-                            mapOf(
-                                previewViewModel.displaySize to
-                                    FullPreviewCropModel(
-                                        it.cropHint,
-                                        it.wallpaperZoom,
-                                        it.hostViewSize,
-                                        it.cropSurfaceSize,
-                                    )
-                            )
-                        )
+                    staticWallpaperPreviewViewModel.run {
+                        updateCropHintsInfo(fullPreviewCropModels)
                     }
                 }
             } else {
diff --git a/src/com/android/wallpaper/util/DisplayUtils.kt b/src/com/android/wallpaper/util/DisplayUtils.kt
index 4763d92..fa1d879 100644
--- a/src/com/android/wallpaper/util/DisplayUtils.kt
+++ b/src/com/android/wallpaper/util/DisplayUtils.kt
@@ -86,6 +86,27 @@
     }
 
     /**
+     * This flag returns true if the display is:
+     * 1. a large screen device display, e.g. tablet
+     * 2. an unfolded display from a foldable device
+     *
+     * This flag returns false the display is:
+     * 1. a handheld device display
+     * 2. a folded display from a foldable device
+     */
+    fun isLargeScreenOrUnfoldedDisplay(activity: Activity): Boolean {
+        // Note that a foldable is a large screen device if the largest display is large screen.
+        // Ths flag is true if it is a large screen device, e.g. tablet, or a foldable device.
+        val isLargeScreenOrFoldable = isLargeScreenDevice()
+        // For a single display device, this flag is always true.
+        // For a multi-display device, it is only true when the current display is the largest
+        // display. For the case of foldable, it is true when the display is the unfolded one, and
+        // false when it is folded.
+        val isSingleDisplayOrUnfolded = isOnWallpaperDisplay(activity)
+        return isLargeScreenOrFoldable && isSingleDisplayOrUnfolded
+    }
+
+    /**
      * Returns true if this device's screen (or largest screen in case of multiple screen devices)
      * is considered a "Large screen"
      */
diff --git a/tests/Android.bp b/tests/Android.bp
index 238d8a2..1004eae 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -58,6 +58,7 @@
         "junit",
         "kotlinx_coroutines_test",
         "truth",
+        "flag-junit",
     ],
     libs: [
         "android.test.runner",
diff --git a/tests/common/src/com/android/wallpaper/testing/FakeWallpaperClient.kt b/tests/common/src/com/android/wallpaper/testing/FakeWallpaperClient.kt
index 7999a5c..58d83c6 100644
--- a/tests/common/src/com/android/wallpaper/testing/FakeWallpaperClient.kt
+++ b/tests/common/src/com/android/wallpaper/testing/FakeWallpaperClient.kt
@@ -17,6 +17,7 @@
 
 package com.android.wallpaper.testing
 
+import android.app.WallpaperColors
 import android.graphics.Bitmap
 import android.graphics.Point
 import android.graphics.Rect
@@ -142,6 +143,13 @@
         return emptyMap()
     }
 
+    override suspend fun getWallpaperColors(
+        bitmap: Bitmap,
+        cropHints: Map<Point, Rect>?
+    ): WallpaperColors? {
+        return null
+    }
+
     companion object {
         val INITIAL_RECENT_WALLPAPERS =
             listOf(
diff --git a/tests/src/com/android/wallpaper/picker/preview/ui/WallpaperPreviewActivityTest.kt b/tests/src/com/android/wallpaper/picker/preview/ui/WallpaperPreviewActivityTest.kt
index 91459de..7778b91 100644
--- a/tests/src/com/android/wallpaper/picker/preview/ui/WallpaperPreviewActivityTest.kt
+++ b/tests/src/com/android/wallpaper/picker/preview/ui/WallpaperPreviewActivityTest.kt
@@ -15,20 +15,25 @@
  */
 package com.android.wallpaper.picker.preview.ui
 
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
 import androidx.navigation.fragment.NavHostFragment
 import androidx.test.core.app.ActivityScenario
 import androidx.test.core.app.ApplicationProvider
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.MediumTest
+import com.android.wallpaper.Flags.FLAG_MULTI_CROP_PREVIEW_UI_FLAG
 import com.android.wallpaper.model.WallpaperInfo
 import com.android.wallpaper.module.InjectorProvider
 import com.android.wallpaper.testing.TestInjector
 import com.android.wallpaper.testing.TestStaticWallpaperInfo
+import com.android.window.flags.Flags.FLAG_MULTI_CROP
 import com.google.common.truth.Truth.assertThat
 import dagger.hilt.android.testing.HiltAndroidRule
 import dagger.hilt.android.testing.HiltAndroidTest
 import javax.inject.Inject
 import org.junit.Before
+import org.junit.Ignore
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -39,6 +44,8 @@
 class WallpaperPreviewActivityTest {
     @get:Rule var hiltRule = HiltAndroidRule(this)
 
+    @get:Rule val setFlagsRule = SetFlagsRule()
+
     @Inject lateinit var testInjector: TestInjector
 
     private val testStaticWallpaper =
@@ -59,6 +66,8 @@
     }
 
     @Test
+    @Ignore("b/327241549")
+    @EnableFlags(FLAG_MULTI_CROP_PREVIEW_UI_FLAG, FLAG_MULTI_CROP)
     fun showsNavHostFragment() {
         val scenario: ActivityScenario<WallpaperPreviewActivity> =
             ActivityScenario.launch(activityStartIntent)