Tabbed navigation in Wallpaper Picker (1/3).

Introduces tabbed navigation to switch between home and lock screen in
wallpaper picker.

Fix: 262924054
Test: Includes new unit tests for new view-model
Test: Manually verified that switching back and forth keeps the previews
alive and selectively switches the bottom few sections in a nice,
animated way with minimal/no jank.

Change-Id: I087b4c03ac29236be0ed7d512b15e6af887c449f
diff --git a/src/com/android/wallpaper/model/CustomizationSectionController.java b/src/com/android/wallpaper/model/CustomizationSectionController.java
index fd2f3d5..58ef178 100644
--- a/src/com/android/wallpaper/model/CustomizationSectionController.java
+++ b/src/com/android/wallpaper/model/CustomizationSectionController.java
@@ -45,6 +45,16 @@
     /**
      * Returns a newly created {@link SectionView} for the section.
      *
+     * @param context The {@link Context} to inflate view.
+     * @param isOnLockScreen Whether we are on the lock screen.
+     */
+    default T createView(Context context, boolean isOnLockScreen) {
+        return createView(context);
+    }
+
+    /**
+     * Returns a newly created {@link SectionView} for the section.
+     *
      * @param context the {@link Context} to inflate view
      */
     T createView(Context context);
@@ -57,4 +67,7 @@
 
     /** Gets called when the section gets transitioned out. */
     default void onTransitionOut() {}
+
+    /** Notifies when the screen was switched. */
+    default void onScreenSwitched(boolean isOnLockScreen) {}
 }
diff --git a/src/com/android/wallpaper/module/CustomizationSections.java b/src/com/android/wallpaper/module/CustomizationSections.java
index 551d82d..2723265 100644
--- a/src/com/android/wallpaper/module/CustomizationSections.java
+++ b/src/com/android/wallpaper/module/CustomizationSections.java
@@ -18,6 +18,30 @@
 /** Interface for carry {@link CustomizationSectionController}s. */
 public interface CustomizationSections {
 
+    /** Enumerates all screens supported by {@code getSectionControllersForScreen}. */
+    enum Screen {
+        LOCK_SCREEN,
+        HOME_SCREEN,
+    }
+
+    /**
+     * Gets a new instance of the section controller list for the given {@link Screen}.
+     *
+     * Note that the section views will be displayed by the list ordering.
+     *
+     * <p>Don't keep the section controllers as singleton since they contain views.
+     */
+    List<CustomizationSectionController<?>> getSectionControllersForScreen(
+            Screen screen,
+            FragmentActivity activity,
+            LifecycleOwner lifecycleOwner,
+            WallpaperColorsViewModel wallpaperColorsViewModel,
+            WorkspaceViewModel workspaceViewModel,
+            PermissionRequester permissionRequester,
+            WallpaperPreviewNavigator wallpaperPreviewNavigator,
+            CustomizationSectionNavigationController sectionNavigationController,
+            @Nullable Bundle savedInstanceState);
+
     /**
      * Gets a new instance of the section controller list.
      *
diff --git a/src/com/android/wallpaper/module/WallpaperPickerSections.java b/src/com/android/wallpaper/module/WallpaperPickerSections.java
index 39daa3d..3078242 100644
--- a/src/com/android/wallpaper/module/WallpaperPickerSections.java
+++ b/src/com/android/wallpaper/module/WallpaperPickerSections.java
@@ -21,6 +21,33 @@
 public final class WallpaperPickerSections implements CustomizationSections {
 
     @Override
+    public List<CustomizationSectionController<?>> getSectionControllersForScreen(
+            Screen screen,
+            FragmentActivity activity,
+            LifecycleOwner lifecycleOwner,
+            WallpaperColorsViewModel wallpaperColorsViewModel,
+            WorkspaceViewModel workspaceViewModel,
+            PermissionRequester permissionRequester,
+            WallpaperPreviewNavigator wallpaperPreviewNavigator,
+            CustomizationSectionNavigationController sectionNavigationController,
+            @Nullable Bundle savedInstanceState) {
+        List<CustomizationSectionController<?>> sectionControllers = new ArrayList<>();
+
+        sectionControllers.add(
+                new WallpaperSectionController(
+                    activity,
+                    lifecycleOwner,
+                    permissionRequester,
+                    wallpaperColorsViewModel,
+                    workspaceViewModel,
+                    sectionNavigationController,
+                    wallpaperPreviewNavigator,
+                    savedInstanceState));
+
+        return sectionControllers;
+    }
+
+    @Override
     public List<CustomizationSectionController<?>> getAllSectionControllers(
             FragmentActivity activity,
             LifecycleOwner lifecycleOwner,
diff --git a/src/com/android/wallpaper/picker/CustomizationPickerFragment.java b/src/com/android/wallpaper/picker/CustomizationPickerFragment.java
index 125ac8b..9757232 100644
--- a/src/com/android/wallpaper/picker/CustomizationPickerFragment.java
+++ b/src/com/android/wallpaper/picker/CustomizationPickerFragment.java
@@ -40,6 +40,8 @@
 import com.android.wallpaper.module.FragmentFactory;
 import com.android.wallpaper.module.Injector;
 import com.android.wallpaper.module.InjectorProvider;
+import com.android.wallpaper.picker.ui.binder.CustomizationPickerBinder;
+import com.android.wallpaper.picker.ui.viewmodel.CustomizationPickerViewModel;
 import com.android.wallpaper.util.ActivityUtils;
 
 import java.util.ArrayList;
@@ -58,6 +60,7 @@
     private NestedScrollView mNestedScrollView;
     @Nullable private Bundle mBackStackSavedInstanceState;
     private final FragmentFactory mFragmentFactory;
+    @Nullable private CustomizationPickerViewModel mViewModel;
 
     public CustomizationPickerFragment() {
         mFragmentFactory = InjectorProvider.getInjector().getFragmentFactory();
@@ -68,7 +71,31 @@
             @Nullable Bundle savedInstanceState) {
         final View view = inflater.inflate(R.layout.collapsing_toolbar_container_layout,
                 container, /* attachToRoot= */ false);
-        setContentView(view, R.layout.fragment_customization_picker);
+        final boolean isUseRevampedUi =
+                InjectorProvider.getInjector().getFlags().isUseRevampedUi(requireContext());
+        if (isUseRevampedUi) {
+            setContentView(view, R.layout.fragment_tabbed_customization_picker);
+            mViewModel = new ViewModelProvider(
+                    this,
+                    CustomizationPickerViewModel.newFactory(
+                            this,
+                            savedInstanceState)
+            ).get(CustomizationPickerViewModel.class);
+
+            final Bundle finalSavedInstanceState = savedInstanceState;
+            CustomizationPickerBinder.bind(
+                    view,
+                    mViewModel,
+                    this,
+                    isOnLockScreen -> getSectionControllers(
+                            isOnLockScreen
+                                    ? CustomizationSections.Screen.LOCK_SCREEN
+                                    : CustomizationSections.Screen.HOME_SCREEN,
+                            finalSavedInstanceState));
+        } else {
+            setContentView(view, R.layout.fragment_customization_picker);
+        }
+
         if (ActivityUtils.isLaunchedFromSettingsRelated(getActivity().getIntent())) {
             setUpToolbar(view, !ActivityEmbeddingUtils.shouldHideNavigateUpButton(
                     getActivity(), /* isSecondLayerPage= */ true));
@@ -76,39 +103,42 @@
             setUpToolbar(view, /* upArrow= */ false);
         }
 
-        ViewGroup sectionContainer = view.findViewById(R.id.section_container);
-        sectionContainer.setOnApplyWindowInsetsListener((v, windowInsets) -> {
-            v.setPadding(
-                    v.getPaddingLeft(),
-                    v.getPaddingTop(),
-                    v.getPaddingRight(),
-                    windowInsets.getSystemWindowInsetBottom());
-            return windowInsets.consumeSystemWindowInsets();
-        });
-        mNestedScrollView = view.findViewById(R.id.scroll_container);
-
         if (mBackStackSavedInstanceState != null) {
             savedInstanceState = mBackStackSavedInstanceState;
             mBackStackSavedInstanceState = null;
         }
 
-        initSections(savedInstanceState);
-        mSectionControllers.forEach(controller ->
-                mNestedScrollView.post(() -> {
-                            final Context context = getContext();
-                            if (context == null) {
-                                Log.w(TAG, "Adding section views with null context");
-                                return;
+        mNestedScrollView = view.findViewById(R.id.scroll_container);
+
+        if (!isUseRevampedUi) {
+            ViewGroup sectionContainer = view.findViewById(R.id.section_container);
+            sectionContainer.setOnApplyWindowInsetsListener((v, windowInsets) -> {
+                v.setPadding(
+                        v.getPaddingLeft(),
+                        v.getPaddingTop(),
+                        v.getPaddingRight(),
+                        windowInsets.getSystemWindowInsetBottom());
+                return windowInsets.consumeSystemWindowInsets();
+            });
+
+            initSections(savedInstanceState);
+            mSectionControllers.forEach(controller ->
+                    mNestedScrollView.post(() -> {
+                                final Context context = getContext();
+                                if (context == null) {
+                                    Log.w(TAG, "Adding section views with null context");
+                                    return;
+                                }
+                                sectionContainer.addView(controller.createView(context));
                             }
-                            sectionContainer.addView(controller.createView(context));
-                        }
-                )
-        );
-        final Bundle savedInstanceStateRef = savedInstanceState;
-        // Post it to the end of adding views to ensure restoring view state the last task.
-        mNestedScrollView.post(() ->
-                restoreViewState(savedInstanceStateRef)
-        );
+                    )
+            );
+
+            final Bundle savedInstanceStateRef = savedInstanceState;
+            // Post it to the end of adding views to ensure restoring view state the last task.
+            view.post(() -> restoreViewState(savedInstanceStateRef));
+        }
+
         return view;
     }
 
@@ -204,6 +234,20 @@
         mSectionControllers.forEach(CustomizationSectionController::release);
         mSectionControllers.clear();
 
+        mSectionControllers.addAll(
+                getAvailableSections(getAvailableSectionControllers(savedInstanceState)));
+    }
+
+    private List<CustomizationSectionController<?>> getAvailableSectionControllers(
+            @Nullable Bundle savedInstanceState) {
+        return getSectionControllers(
+                null,
+                savedInstanceState);
+    }
+
+    private List<CustomizationSectionController<?>> getSectionControllers(
+            @Nullable CustomizationSections.Screen screen,
+            @Nullable Bundle savedInstanceState) {
         final Injector injector = InjectorProvider.getInjector();
 
         WallpaperColorsViewModel wcViewModel = new ViewModelProvider(getActivity())
@@ -212,18 +256,28 @@
                 .get(WorkspaceViewModel.class);
 
         CustomizationSections sections = injector.getCustomizationSections(getActivity());
-        List<CustomizationSectionController<?>> allSectionControllers =
-                sections.getAllSectionControllers(
-                        getActivity(),
-                        getViewLifecycleOwner(),
-                        wcViewModel,
-                        workspaceViewModel,
-                        getPermissionRequester(),
-                        getWallpaperPreviewNavigator(),
-                        this,
-                        savedInstanceState);
-
-        mSectionControllers.addAll(getAvailableSections(allSectionControllers));
+        if (screen == null) {
+            return sections.getAllSectionControllers(
+                    getActivity(),
+                    getViewLifecycleOwner(),
+                    wcViewModel,
+                    workspaceViewModel,
+                    getPermissionRequester(),
+                    getWallpaperPreviewNavigator(),
+                    this,
+                    savedInstanceState);
+        } else {
+            return sections.getSectionControllersForScreen(
+                    screen,
+                    getActivity(),
+                    getViewLifecycleOwner(),
+                    wcViewModel,
+                    workspaceViewModel,
+                    getPermissionRequester(),
+                    getWallpaperPreviewNavigator(),
+                    this,
+                    savedInstanceState);
+        }
     }
 
     protected List<CustomizationSectionController<?>> getAvailableSections(
diff --git a/src/com/android/wallpaper/picker/ui/binder/CustomizationPickerBinder.kt b/src/com/android/wallpaper/picker/ui/binder/CustomizationPickerBinder.kt
new file mode 100644
index 0000000..ee70538
--- /dev/null
+++ b/src/com/android/wallpaper/picker/ui/binder/CustomizationPickerBinder.kt
@@ -0,0 +1,131 @@
+/*
+ * 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.wallpaper.picker.ui.binder
+
+import android.view.View
+import android.view.ViewGroup
+import android.view.WindowInsets
+import android.widget.FrameLayout
+import androidx.core.view.children
+import androidx.core.view.updateLayoutParams
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleOwner
+import androidx.lifecycle.lifecycleScope
+import androidx.lifecycle.repeatOnLifecycle
+import com.android.wallpaper.R
+import com.android.wallpaper.model.CustomizationSectionController
+import com.android.wallpaper.model.WallpaperSectionController
+import com.android.wallpaper.picker.SectionView
+import com.android.wallpaper.picker.ui.viewmodel.CustomizationPickerViewModel
+import kotlinx.coroutines.launch
+
+typealias SectionController = CustomizationSectionController<*>
+
+/** Binds view to view-model for the customization picker. */
+object CustomizationPickerBinder {
+    @JvmStatic
+    fun bind(
+        view: View,
+        viewModel: CustomizationPickerViewModel,
+        lifecycleOwner: LifecycleOwner,
+        sectionControllerProvider: (isOnLockScreen: Boolean) -> List<SectionController>,
+    ) {
+        CustomizationPickerTabsBinder.bind(
+            view = view,
+            viewModel = viewModel,
+            lifecycleOwner = lifecycleOwner,
+        )
+
+        val sectionContainer = view.findViewById<ViewGroup>(R.id.section_container)
+        sectionContainer.setOnApplyWindowInsetsListener { v: View, windowInsets: WindowInsets ->
+            v.setPadding(
+                v.paddingLeft,
+                v.paddingTop,
+                v.paddingRight,
+                windowInsets.systemWindowInsetBottom
+            )
+            windowInsets.consumeSystemWindowInsets()
+        }
+        sectionContainer.updateLayoutParams<FrameLayout.LayoutParams> {
+            // We don't want the top margin from the XML because our tabs have that as padding so
+            // they can be collapsed into the toolbar with spacing from the actual title text.
+            topMargin = 0
+        }
+
+        lifecycleOwner.lifecycleScope.launch {
+            lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+                launch {
+                    viewModel.isOnLockScreen.collect { isOnLockScreen ->
+                        // These are the available section controllers we should use now.
+                        val newSectionControllers =
+                            sectionControllerProvider.invoke(isOnLockScreen).filter {
+                                it.isAvailable(view.context)
+                            }
+
+                        check(newSectionControllers[0] is WallpaperSectionController) {
+                            "The first section must always be the wallpaper preview or the" +
+                                " assumption below must be updated."
+                        }
+
+                        val firstTime = sectionContainer.childCount == 0
+                        if (!firstTime) {
+                            // Remove all views, except the very first one, which we assume is for
+                            // the wallpaper preview section.
+                            sectionContainer.removeViews(1, sectionContainer.childCount - 1)
+
+                            // The old controllers for the removed views should be released, except
+                            // for the very first one, which is for the wallpaper preview section;
+                            // that one we keep but just tell it that we switched screens.
+                            sectionContainer.children
+                                .mapNotNull { it.tag as? SectionController }
+                                .forEachIndexed { index, oldController ->
+                                    if (index == 0) {
+                                        // We assume that index 0 is the wallpaper preview section.
+                                        // We keep it because it's an expensive section (as it needs
+                                        // to maintain a wallpaper connection that seems to be
+                                        // making assumptions about its SurfaceView always remaining
+                                        // attached to the window).
+                                        oldController.onScreenSwitched(isOnLockScreen)
+                                    } else {
+                                        // All other old controllers will be thrown out so let's
+                                        // release them.
+                                        oldController.release()
+                                    }
+                                }
+                        }
+
+                        // Let's add the new controllers and views.
+                        newSectionControllers.forEachIndexed { index, controller ->
+                            if (firstTime || index > 0) {
+                                val addedView = controller.createView(view.context, isOnLockScreen)
+                                addedView.tag = controller
+                                sectionContainer.addView(addedView)
+                            }
+                        }
+                    }
+                }
+            }
+
+            // This happens when the lifecycle is stopped.
+            sectionContainer.children
+                .mapNotNull { it.tag as? CustomizationSectionController<out SectionView> }
+                .forEach { controller -> controller.release() }
+            sectionContainer.removeAllViews()
+        }
+    }
+}
diff --git a/src/com/android/wallpaper/picker/ui/binder/CustomizationPickerTabsBinder.kt b/src/com/android/wallpaper/picker/ui/binder/CustomizationPickerTabsBinder.kt
new file mode 100644
index 0000000..dfd8ca6
--- /dev/null
+++ b/src/com/android/wallpaper/picker/ui/binder/CustomizationPickerTabsBinder.kt
@@ -0,0 +1,74 @@
+/*
+ * 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.wallpaper.picker.ui.binder
+
+import android.view.View
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleOwner
+import androidx.lifecycle.lifecycleScope
+import androidx.lifecycle.repeatOnLifecycle
+import com.android.wallpaper.R
+import com.android.wallpaper.picker.ui.viewmodel.CustomizationPickerViewModel
+import com.android.wallpaper.widget.DuoTabs
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.launch
+
+/** Binds view to view-model for the customization picker tabs. */
+object CustomizationPickerTabsBinder {
+    @JvmStatic
+    fun bind(
+        view: View,
+        viewModel: CustomizationPickerViewModel,
+        lifecycleOwner: LifecycleOwner,
+    ) {
+        val tabs: DuoTabs = view.requireViewById(R.id.duo_tabs)
+        tabs.setTabText(
+            view.context.getString(R.string.lock_screen_tab),
+            view.context.getString(R.string.home_screen_tab),
+        )
+
+        lifecycleOwner.lifecycleScope.launch {
+            lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+                launch {
+                    combine(viewModel.lockScreenTab, viewModel.homeScreenTab) {
+                            lockScreenTabViewModel,
+                            homeScreenTabViewModel ->
+                            lockScreenTabViewModel to homeScreenTabViewModel
+                        }
+                        .collect { (lockScreenTabViewModel, homeScreenTabViewModel) ->
+                            tabs.setOnTabSelectedListener(null)
+                            tabs.selectTab(
+                                when {
+                                    lockScreenTabViewModel.isSelected -> DuoTabs.TAB_PRIMARY
+                                    else -> DuoTabs.TAB_SECONDARY
+                                },
+                            )
+                            tabs.setOnTabSelectedListener { tabId ->
+                                when (tabId) {
+                                    DuoTabs.TAB_PRIMARY ->
+                                        lockScreenTabViewModel.onClicked?.invoke()
+                                    DuoTabs.TAB_SECONDARY ->
+                                        homeScreenTabViewModel.onClicked?.invoke()
+                                }
+                            }
+                        }
+                }
+            }
+        }
+    }
+}
diff --git a/src/com/android/wallpaper/picker/ui/viewmodel/CustomizationPickerTabViewModel.kt b/src/com/android/wallpaper/picker/ui/viewmodel/CustomizationPickerTabViewModel.kt
new file mode 100644
index 0000000..32cf4c8
--- /dev/null
+++ b/src/com/android/wallpaper/picker/ui/viewmodel/CustomizationPickerTabViewModel.kt
@@ -0,0 +1,24 @@
+/*
+ * 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.wallpaper.picker.ui.viewmodel
+
+/** Models UI state for a single tab. */
+data class CustomizationPickerTabViewModel(
+    val isSelected: Boolean,
+    val onClicked: (() -> Unit)?,
+)
diff --git a/src/com/android/wallpaper/picker/ui/viewmodel/CustomizationPickerViewModel.kt b/src/com/android/wallpaper/picker/ui/viewmodel/CustomizationPickerViewModel.kt
new file mode 100644
index 0000000..0a3b7b9
--- /dev/null
+++ b/src/com/android/wallpaper/picker/ui/viewmodel/CustomizationPickerViewModel.kt
@@ -0,0 +1,98 @@
+/*
+ * 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.wallpaper.picker.ui.viewmodel
+
+import android.os.Bundle
+import androidx.annotation.VisibleForTesting
+import androidx.lifecycle.AbstractSavedStateViewModelFactory
+import androidx.lifecycle.SavedStateHandle
+import androidx.lifecycle.ViewModel
+import androidx.savedstate.SavedStateRegistryOwner
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.map
+
+/** Models UI state for the customization picker. */
+class CustomizationPickerViewModel
+@VisibleForTesting
+constructor(
+    private val savedStateHandle: SavedStateHandle,
+) : ViewModel() {
+
+    private val _isOnLockScreen = MutableStateFlow(true)
+    /** Whether we are on the lock screen. If `false`, we are on the home screen. */
+    val isOnLockScreen: Flow<Boolean> = _isOnLockScreen.asStateFlow()
+
+    /** A view-model for the "lock screen" tab. */
+    val lockScreenTab: Flow<CustomizationPickerTabViewModel> =
+        isOnLockScreen.map { onLockScreen ->
+            CustomizationPickerTabViewModel(
+                isSelected = onLockScreen,
+                onClicked =
+                    if (!onLockScreen) {
+                        {
+                            _isOnLockScreen.value = true
+                            savedStateHandle[KEY_SAVED_STATE_IS_ON_LOCK_SCREEN] = true
+                        }
+                    } else {
+                        null
+                    }
+            )
+        }
+    /** A view-model for the "home screen" tab. */
+    val homeScreenTab: Flow<CustomizationPickerTabViewModel> =
+        isOnLockScreen.map { onLockScreen ->
+            CustomizationPickerTabViewModel(
+                isSelected = !onLockScreen,
+                onClicked =
+                    if (onLockScreen) {
+                        {
+                            _isOnLockScreen.value = false
+                            savedStateHandle[KEY_SAVED_STATE_IS_ON_LOCK_SCREEN] = false
+                        }
+                    } else {
+                        null
+                    }
+            )
+        }
+
+    init {
+        _isOnLockScreen.value = savedStateHandle[KEY_SAVED_STATE_IS_ON_LOCK_SCREEN] ?: true
+    }
+
+    companion object {
+        @JvmStatic
+        fun newFactory(
+            owner: SavedStateRegistryOwner,
+            defaultArgs: Bundle? = null,
+        ): AbstractSavedStateViewModelFactory =
+            object : AbstractSavedStateViewModelFactory(owner, defaultArgs) {
+                @Suppress("UNCHECKED_CAST")
+                override fun <T : ViewModel> create(
+                    key: String,
+                    modelClass: Class<T>,
+                    handle: SavedStateHandle,
+                ): T {
+                    return CustomizationPickerViewModel(handle) as T
+                }
+            }
+
+        private const val KEY_SAVED_STATE_IS_ON_LOCK_SCREEN = "is_on_lock_screen"
+    }
+}