Fix preview screen issue on foldables when unfolding

Test: Manually tested the issue was gone
Bug: 326126791
Flag: ACONFIG com.android.wallpaper.multi_crop_preview_ui_flag DEVELOPMENT
Change-Id: I62ac97adb750020336055bfac581d5d2ec816cf9
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/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/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"
      */