Merge "Adds Google account login during private space setup" into main
diff --git a/res/drawable/ic_error_red.xml b/res/drawable/ic_error_red.xml
new file mode 100644
index 0000000..d17c85b
--- /dev/null
+++ b/res/drawable/ic_error_red.xml
@@ -0,0 +1,26 @@
+<!--
+  ~ 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.
+  -->
+
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+        android:width="24dp"
+        android:height="24dp"
+        android:viewportWidth="960"
+        android:viewportHeight="960"
+        android:tint="?android:attr/colorError">
+    <path
+        android:fillColor="@android:color/white"
+        android:pathData="M480,680Q497,680 508.5,668.5Q520,657 520,640Q520,623 508.5,611.5Q497,600 480,600Q463,600 451.5,611.5Q440,623 440,640Q440,657 451.5,668.5Q463,680 480,680ZM440,520L520,520L520,280L440,280L440,520ZM480,880Q397,880 324,848.5Q251,817 197,763Q143,709 111.5,636Q80,563 80,480Q80,397 111.5,324Q143,251 197,197Q251,143 324,111.5Q397,80 480,80Q563,80 636,111.5Q709,143 763,197Q817,251 848.5,324Q880,397 880,480Q880,563 848.5,636Q817,709 763,763Q709,817 636,848.5Q563,880 480,880ZM480,800Q614,800 707,707Q800,614 800,480Q800,346 707,253Q614,160 480,160Q346,160 253,253Q160,346 160,480Q160,614 253,707Q346,800 480,800ZM480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Z"/>
+</vector>
diff --git a/res/layout/privatespace_account_login_error.xml b/res/layout/privatespace_account_login_error.xml
new file mode 100644
index 0000000..a38dd50
--- /dev/null
+++ b/res/layout/privatespace_account_login_error.xml
@@ -0,0 +1,28 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ 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.
+  -->
+
+<com.google.android.setupdesign.GlifLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:app="http://schemas.android.com/apk/res-auto"
+    android:id="@+id/ps_error_page_layout"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent"
+    android:filterTouchesWhenObscured="true"
+    app:sucHeaderText="@string/privatespace_retry_signin_title"
+    app:sudDescriptionText="@string/privatespace_retry_summary"
+    android:icon="@drawable/ic_error_red">
+</com.google.android.setupdesign.GlifLayout>
diff --git a/res/navigation/privatespace_main_context_nav.xml b/res/navigation/privatespace_main_context_nav.xml
index 5b2552a..ffc63ec 100644
--- a/res/navigation/privatespace_main_context_nav.xml
+++ b/res/navigation/privatespace_main_context_nav.xml
@@ -33,8 +33,8 @@
             android:id="@+id/action_advance_profile_error"
             app:destination="@id/ps_profile_error_fragment"/>
         <action
-            android:id="@+id/action_advance_to_success"
-            app:destination="@id/ps_profile_success_fragment"/>
+            android:id="@+id/action_advance_login_error"
+            app:destination="@id/ps_account_error_fragment"/>
     </fragment>
     <fragment android:id="@+id/ps_profile_error_fragment"
               android:name="com.android.settings.privatespace.PrivateProfileCreationError"
@@ -46,4 +46,18 @@
     <fragment android:id="@+id/ps_profile_success_fragment"
               android:name="com.android.settings.privatespace.SetupSuccessFragment"
               android:label="fragment_ps_success"/>
-</navigation>
\ No newline at end of file
+    <fragment android:id="@+id/ps_account_error_fragment"
+              android:name="com.android.settings.privatespace.PrivateSpaceAccountLoginError"
+              android:label="fragment_account_error">
+        <action
+            android:id="@+id/action_advance_login_error"
+            app:destination="@id/ps_account_error_fragment"/>
+    </fragment>
+    <fragment android:id="@+id/ps_profile_lock_fragment"
+              android:name="com.android.settings.privatespace.PrivateSpaceSetLockFragment"
+              android:label="fragment_ps_lock"/>
+    <action android:id="@+id/action_success_fragment"
+            app:destination="@id/ps_profile_success_fragment"/>
+    <action android:id="@+id/action_set_lock_fragment"
+            app:destination="@id/ps_profile_lock_fragment"/>
+</navigation>
diff --git a/res/navigation/privatespace_private_context_nav.xml b/res/navigation/privatespace_private_context_nav.xml
deleted file mode 100644
index 3df8fa5..0000000
--- a/res/navigation/privatespace_private_context_nav.xml
+++ /dev/null
@@ -1,24 +0,0 @@
-<!--
-  ~ 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.
-  -->
-
-<navigation xmlns:android="http://schemas.android.com/apk/res/android"
-            xmlns:app="http://schemas.android.com/apk/res-auto"
-            android:id="@+id/private_space_navigation"
-            app:startDestination="@id/ps_profile_lock_fragment">
-    <fragment android:id="@+id/ps_profile_lock_fragment"
-              android:name="com.android.settings.privatespace.PrivateSpaceSetLockFragment"
-              android:label="fragment_ps_lock"/>
-</navigation>
diff --git a/res/values/config.xml b/res/values/config.xml
index 99052ca..f50e918 100755
--- a/res/values/config.xml
+++ b/res/values/config.xml
@@ -753,6 +753,9 @@
     <!-- Whether to display Cloned Apps page in Settings (Settings > Apps > Cloned Apps).-->
     <bool name="config_cloned_apps_page_enabled">false</bool>
 
+    <!-- Whether to initiate Account login during Private Space setup.-->
+    <bool name="config_privatespace_account_login_enabled">false</bool>
+
     <!-- Certificates of apps which are allowed to use activity embedding with Settings.-->
     <string-array name="config_known_host_certs" translatable="false">
         <item></item>
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 704e500..9d7ca65 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -1298,6 +1298,10 @@
     <string name="privatespace_done_label">Done</string>
     <!-- Toast to show on private space setup completion informing user to scroll down All apps to access private space. [CHAR LIMIT=60] -->
     <string name="scrolldown_to_access">Scroll down to access Private Space</string>
+    <!-- Title for Private Space account login error screen. [CHAR LIMIT=60] -->
+    <string name="privatespace_retry_signin_title">Sign in to set up Private Space</string>
+    <!-- Summary for the Private Space account login error screen. [CHAR LIMIT=NONE] -->
+    <string name="privatespace_retry_summary">You need to sign in to a Account to set up Private Space</string>
 
     <!-- Text shown when "Add fingerprint" button is disabled -->
     <string name="fingerprint_add_max">You can add up to <xliff:g id="count" example="5">%d</xliff:g> fingerprints</string>
diff --git a/src/com/android/settings/overlay/FeatureFactory.kt b/src/com/android/settings/overlay/FeatureFactory.kt
index ac689d9..bc0cf1f 100644
--- a/src/com/android/settings/overlay/FeatureFactory.kt
+++ b/src/com/android/settings/overlay/FeatureFactory.kt
@@ -39,6 +39,7 @@
 import com.android.settings.onboarding.OnboardingFeatureProvider
 import com.android.settings.overlay.FeatureFactory.Companion.setFactory
 import com.android.settings.panel.PanelFeatureProvider
+import com.android.settings.privatespace.PrivateSpaceLoginFeatureProvider
 import com.android.settings.search.SearchFeatureProvider
 import com.android.settings.security.SecurityFeatureProvider
 import com.android.settings.security.SecuritySettingsFeatureProvider
@@ -170,6 +171,11 @@
      */
     abstract val fastPairFeatureProvider: FastPairFeatureProvider
 
+    /**
+     * Gets implementation for Private Space account login feature.
+     */
+    abstract val privateSpaceLoginFeatureProvider: PrivateSpaceLoginFeatureProvider
+
     companion object {
         private var _factory: FeatureFactory? = null
 
diff --git a/src/com/android/settings/overlay/FeatureFactoryImpl.kt b/src/com/android/settings/overlay/FeatureFactoryImpl.kt
index 7f991b7..28dbb23 100644
--- a/src/com/android/settings/overlay/FeatureFactoryImpl.kt
+++ b/src/com/android/settings/overlay/FeatureFactoryImpl.kt
@@ -59,6 +59,8 @@
 import com.android.settings.security.SecurityFeatureProviderImpl
 import com.android.settings.security.SecuritySettingsFeatureProvider
 import com.android.settings.security.SecuritySettingsFeatureProviderImpl
+import com.android.settings.privatespace.PrivateSpaceLoginFeatureProvider
+import com.android.settings.privatespace.PrivateSpaceLoginFeatureProviderImpl
 import com.android.settings.slices.SlicesFeatureProviderImpl
 import com.android.settings.users.UserFeatureProviderImpl
 import com.android.settings.vpn2.AdvancedVpnFeatureProviderImpl
@@ -184,4 +186,8 @@
     override val fastPairFeatureProvider: FastPairFeatureProvider by lazy {
         FastPairFeatureProviderImpl()
     }
+
+    override val privateSpaceLoginFeatureProvider: PrivateSpaceLoginFeatureProvider by lazy {
+        PrivateSpaceLoginFeatureProviderImpl()
+    }
 }
diff --git a/src/com/android/settings/privatespace/AutoAdvanceSetupFragment.java b/src/com/android/settings/privatespace/AutoAdvanceSetupFragment.java
index 5456c01..3b59166 100644
--- a/src/com/android/settings/privatespace/AutoAdvanceSetupFragment.java
+++ b/src/com/android/settings/privatespace/AutoAdvanceSetupFragment.java
@@ -16,7 +16,8 @@
 
 package com.android.settings.privatespace;
 
-import static com.android.settings.privatespace.PrivateSpaceSetupActivity.SET_LOCK_ACTION;
+import static com.android.settings.privatespace.PrivateSpaceSetupActivity.ACCOUNT_LOGIN_ACTION;
+import static com.android.settings.privatespace.PrivateSpaceSetupActivity.EXTRA_ACTION_TYPE;
 
 import android.annotation.SuppressLint;
 import android.content.Intent;
@@ -134,7 +135,8 @@
     private void startActivityInPrivateUser(UserHandle userHandle) {
         /* Start new activity in private profile which is needed to set private profile lock */
         Intent intent = new Intent(getContext(), PrivateProfileContextHelperActivity.class);
-        getActivity().startActivityForResultAsUser(intent, SET_LOCK_ACTION, userHandle);
+        intent.putExtra(EXTRA_ACTION_TYPE, ACCOUNT_LOGIN_ACTION);
+        getActivity().startActivityForResultAsUser(intent, ACCOUNT_LOGIN_ACTION, userHandle);
     }
 
     private void showPrivateSpaceErrorScreen() {
diff --git a/src/com/android/settings/privatespace/PrivateProfileContextHelperActivity.java b/src/com/android/settings/privatespace/PrivateProfileContextHelperActivity.java
index c0d762a..0539f60 100644
--- a/src/com/android/settings/privatespace/PrivateProfileContextHelperActivity.java
+++ b/src/com/android/settings/privatespace/PrivateProfileContextHelperActivity.java
@@ -16,27 +16,84 @@
 
 package com.android.settings.privatespace;
 
+import static android.app.admin.DevicePolicyManager.ACTION_SET_NEW_PASSWORD;
+import static android.app.admin.DevicePolicyManager.EXTRA_PASSWORD_COMPLEXITY;
+import static android.app.admin.DevicePolicyManager.PASSWORD_COMPLEXITY_LOW;
+
+import static com.android.settings.privatespace.PrivateSpaceSetupActivity.ACCOUNT_LOGIN_ACTION;
+import static com.android.settings.privatespace.PrivateSpaceSetupActivity.EXTRA_ACTION_TYPE;
+import static com.android.settings.privatespace.PrivateSpaceSetupActivity.SET_LOCK_ACTION;
+
+import android.app.KeyguardManager;
+import android.content.Intent;
 import android.os.Bundle;
 
+import androidx.activity.result.ActivityResult;
+import androidx.activity.result.ActivityResultLauncher;
+import androidx.activity.result.contract.ActivityResultContracts;
+import androidx.annotation.Nullable;
 import androidx.fragment.app.FragmentActivity;
-import androidx.navigation.fragment.NavHostFragment;
 
-import com.android.settings.R;
 import com.android.settings.SetupWizardUtils;
+import com.android.settings.overlay.FeatureFactory;
 
 import com.google.android.setupdesign.util.ThemeHelper;
 
-/** Activity that is started as private profile user that helps to set private profile lock. */
+/** Activity that is started as private profile user that helps to set private profile lock or
+ * add an account on the private profile. */
 public class PrivateProfileContextHelperActivity extends FragmentActivity {
     private static final String TAG = "PrivateProfileHelper";
+    private final ActivityResultLauncher<Intent> mAddAccountToPrivateProfile =
+            registerForActivityResult(new ActivityResultContracts.StartActivityForResult(),
+                    this::onAccountAdded);
+    private final ActivityResultLauncher<Intent> mVerifyDeviceLock =
+            registerForActivityResult(new ActivityResultContracts.StartActivityForResult(),
+                    this::onSetDeviceNewLock);
+
     @Override
     protected void onCreate(Bundle savedInstanceState) {
         setTheme(SetupWizardUtils.getTheme(this, getIntent()));
         ThemeHelper.trySetDynamicColor(this);
         super.onCreate(savedInstanceState);
-        setContentView(R.layout.privatespace_setup_root);
-        NavHostFragment navHostFragment = (NavHostFragment) getSupportFragmentManager()
-                .findFragmentById(R.id.ps_nav_host_fragment);
-        navHostFragment.getNavController().setGraph(R.navigation.privatespace_private_context_nav);
+        if (savedInstanceState == null) {
+            int action = getIntent().getIntExtra(EXTRA_ACTION_TYPE, -1);
+            if (action == ACCOUNT_LOGIN_ACTION) {
+                PrivateSpaceLoginFeatureProvider privateSpaceLoginFeatureProvider =
+                        FeatureFactory.getFeatureFactory().getPrivateSpaceLoginFeatureProvider();
+                if (!privateSpaceLoginFeatureProvider.initiateAccountLogin(this,
+                        mAddAccountToPrivateProfile)) {
+                    setResult(RESULT_OK);
+                    finish();
+                }
+            } else if (action == SET_LOCK_ACTION) {
+                createPrivateSpaceLock();
+            }
+        }
+    }
+
+    private void createPrivateSpaceLock() {
+        final Intent intent = new Intent(ACTION_SET_NEW_PASSWORD);
+        intent.putExtra(EXTRA_PASSWORD_COMPLEXITY, PASSWORD_COMPLEXITY_LOW);
+        mVerifyDeviceLock.launch(intent);
+    }
+
+    private void onAccountAdded(@Nullable ActivityResult result) {
+        if (result != null && result.getResultCode() == RESULT_OK) {
+            setResult(RESULT_OK);
+        } else {
+            setResult(RESULT_CANCELED);
+        }
+        finish();
+    }
+
+    private void onSetDeviceNewLock(@Nullable ActivityResult result) {
+        // TODO(b/307281644) : Verify this for biometrics and check result code after new
+        //  Authentication changes are merged.
+        if (result != null && getSystemService(KeyguardManager.class).isDeviceSecure()) {
+            setResult(RESULT_OK);
+        } else {
+            setResult(RESULT_CANCELED);
+        }
+        finish();
     }
 }
diff --git a/src/com/android/settings/privatespace/PrivateSpaceAccountLoginError.java b/src/com/android/settings/privatespace/PrivateSpaceAccountLoginError.java
new file mode 100644
index 0000000..2d263d7
--- /dev/null
+++ b/src/com/android/settings/privatespace/PrivateSpaceAccountLoginError.java
@@ -0,0 +1,85 @@
+/*
+ * 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.settings.privatespace;
+
+import static com.android.settings.privatespace.PrivateSpaceSetupActivity.ACCOUNT_LOGIN_ACTION;
+import static com.android.settings.privatespace.PrivateSpaceSetupActivity.EXTRA_ACTION_TYPE;
+
+import android.annotation.SuppressLint;
+import android.content.Intent;
+import android.os.Bundle;
+import android.os.UserHandle;
+import android.view.LayoutInflater;
+import android.view.View;
+import android.view.ViewGroup;
+
+import androidx.activity.OnBackPressedCallback;
+import androidx.annotation.Nullable;
+import androidx.fragment.app.Fragment;
+
+import com.android.settings.R;
+
+import com.google.android.setupcompat.template.FooterBarMixin;
+import com.google.android.setupcompat.template.FooterButton;
+import com.google.android.setupdesign.GlifLayout;
+
+/** Fragment to display error screen if the profile is not signed in with a Google account. */
+public class PrivateSpaceAccountLoginError extends Fragment {
+    @Override
+    public View onCreateView(
+            LayoutInflater inflater,
+            @Nullable ViewGroup container,
+            @Nullable Bundle savedInstanceState) {
+        GlifLayout rootView =
+                (GlifLayout) inflater
+                        .inflate(R.layout.privatespace_account_login_error, container, false);
+        final FooterBarMixin mixin = rootView.getMixin(FooterBarMixin.class);
+        mixin.setPrimaryButton(
+                new FooterButton.Builder(getContext())
+                        .setText(R.string.privatespace_tryagain_label)
+                        .setListener(nextScreen())
+                        .setButtonType(FooterButton.ButtonType.NEXT)
+                        .setTheme(com.google.android.setupdesign.R.style.SudGlifButton_Primary)
+                        .build());
+        OnBackPressedCallback callback =
+                new OnBackPressedCallback(true /* enabled by default */) {
+                    @Override
+                    public void handleOnBackPressed() {
+                        // Handle the back button event
+                    }
+                };
+        requireActivity().getOnBackPressedDispatcher().addCallback(this, callback);
+
+        return rootView;
+    }
+
+    @SuppressLint("MissingPermission")
+    private View.OnClickListener nextScreen() {
+        return v -> {
+            PrivateSpaceMaintainer privateSpaceMaintainer = PrivateSpaceMaintainer
+                    .getInstance(getActivity());
+            UserHandle userHandle;
+            if (privateSpaceMaintainer.doesPrivateSpaceExist() && (userHandle =
+                    privateSpaceMaintainer.getPrivateProfileHandle()) != null) {
+                Intent intent = new Intent(getContext(), PrivateProfileContextHelperActivity.class);
+                intent.putExtra(EXTRA_ACTION_TYPE, ACCOUNT_LOGIN_ACTION);
+                getActivity().startActivityForResultAsUser(intent, ACCOUNT_LOGIN_ACTION,
+                        userHandle);
+            }
+        };
+    }
+}
diff --git a/src/com/android/settings/privatespace/PrivateSpaceLoginFeatureProvider.java b/src/com/android/settings/privatespace/PrivateSpaceLoginFeatureProvider.java
new file mode 100644
index 0000000..76ea9ac
--- /dev/null
+++ b/src/com/android/settings/privatespace/PrivateSpaceLoginFeatureProvider.java
@@ -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.settings.privatespace;
+
+import android.content.Context;
+import android.content.Intent;
+
+import androidx.activity.result.ActivityResultLauncher;
+import androidx.annotation.NonNull;
+
+/** Feature provider for account login during Private Space setup. */
+public interface PrivateSpaceLoginFeatureProvider {
+    /** Returns true if login to an account is enabled during Private Space setup. */
+    boolean initiateAccountLogin(@NonNull Context context,
+            @NonNull ActivityResultLauncher<Intent> resultLauncher);
+}
diff --git a/src/com/android/settings/privatespace/PrivateSpaceLoginFeatureProviderImpl.java b/src/com/android/settings/privatespace/PrivateSpaceLoginFeatureProviderImpl.java
new file mode 100644
index 0000000..7fca2a4
--- /dev/null
+++ b/src/com/android/settings/privatespace/PrivateSpaceLoginFeatureProviderImpl.java
@@ -0,0 +1,32 @@
+/*
+ * 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.settings.privatespace;
+
+import android.content.Context;
+import android.content.Intent;
+
+import androidx.activity.result.ActivityResultLauncher;
+import androidx.annotation.NonNull;
+
+/** Stub class for Private space to not initiate account login during setup */
+public class PrivateSpaceLoginFeatureProviderImpl implements PrivateSpaceLoginFeatureProvider {
+    @Override
+    public boolean initiateAccountLogin(@NonNull Context context,
+            @NonNull ActivityResultLauncher<Intent> resultLauncher) {
+        return false;
+    }
+}
diff --git a/src/com/android/settings/privatespace/PrivateSpaceSetLockFragment.java b/src/com/android/settings/privatespace/PrivateSpaceSetLockFragment.java
index 3d17638..93dc43b 100644
--- a/src/com/android/settings/privatespace/PrivateSpaceSetLockFragment.java
+++ b/src/com/android/settings/privatespace/PrivateSpaceSetLockFragment.java
@@ -16,25 +16,21 @@
 
 package com.android.settings.privatespace;
 
-import static android.app.Activity.RESULT_OK;
-import static android.app.admin.DevicePolicyManager.ACTION_SET_NEW_PASSWORD;
-import static android.app.admin.DevicePolicyManager.EXTRA_PASSWORD_COMPLEXITY;
-import static android.app.admin.DevicePolicyManager.PASSWORD_COMPLEXITY_LOW;
+import static com.android.settings.privatespace.PrivateSpaceSetupActivity.EXTRA_ACTION_TYPE;
+import static com.android.settings.privatespace.PrivateSpaceSetupActivity.SET_LOCK_ACTION;
 
-import android.app.Activity;
-import android.app.KeyguardManager;
+import android.annotation.SuppressLint;
 import android.content.Intent;
 import android.os.Bundle;
+import android.os.UserHandle;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.view.ViewGroup;
 
 import androidx.activity.OnBackPressedCallback;
-import androidx.activity.result.ActivityResult;
-import androidx.activity.result.ActivityResultLauncher;
-import androidx.activity.result.contract.ActivityResultContracts;
 import androidx.annotation.Nullable;
 import androidx.fragment.app.Fragment;
+import androidx.navigation.fragment.NavHostFragment;
 
 import com.android.settings.R;
 
@@ -45,9 +41,6 @@
 /** Fragment that provides an option to user to choose between the existing screen lock or set a
  * separate private profile lock. */
 public class PrivateSpaceSetLockFragment extends Fragment {
-    private final ActivityResultLauncher<Intent> mVerifyDeviceLock =
-            registerForActivityResult(new ActivityResultContracts.StartActivityForResult(),
-                    this::onSetDeviceNewLock);
 
     @Override
     public View onCreateView(
@@ -90,11 +83,8 @@
     private View.OnClickListener onClickUse() {
         return v -> {
             // Simply Use default screen lock. No need to handle
-            Activity activity = getActivity();
-            if (activity != null) {
-                activity.setResult(RESULT_OK);
-                activity.finish();
-            }
+            NavHostFragment.findNavController(PrivateSpaceSetLockFragment.this)
+                    .navigate(R.id.action_success_fragment);
         };
     }
 
@@ -104,22 +94,17 @@
         };
     }
 
+    @SuppressLint("MissingPermission")
     private void createPrivateSpaceLock() {
-        final Intent intent = new Intent(ACTION_SET_NEW_PASSWORD);
-        intent.putExtra(EXTRA_PASSWORD_COMPLEXITY, PASSWORD_COMPLEXITY_LOW);
-        mVerifyDeviceLock.launch(intent);
-    }
-
-    private void onSetDeviceNewLock(@Nullable ActivityResult result) {
-        // TODO(b/307281644) : Verify this for biometrics and check result code after new
-        //  Authentication changes are merged.
-        if (result != null) {
-            Activity profileContextHelperActivity = getActivity();
-            if (profileContextHelperActivity != null && profileContextHelperActivity
-                    .getSystemService(KeyguardManager.class).isDeviceSecure()) {
-                profileContextHelperActivity.setResult(RESULT_OK);
-                profileContextHelperActivity.finish();
-            }
+        PrivateSpaceMaintainer privateSpaceMaintainer = PrivateSpaceMaintainer
+                .getInstance(getActivity());
+        UserHandle userHandle;
+        if (privateSpaceMaintainer.doesPrivateSpaceExist() && (userHandle =
+                privateSpaceMaintainer.getPrivateProfileHandle()) != null) {
+            Intent intent = new Intent(getContext(), PrivateProfileContextHelperActivity.class);
+            intent.putExtra(EXTRA_ACTION_TYPE, SET_LOCK_ACTION);
+            getActivity().startActivityForResultAsUser(intent, SET_LOCK_ACTION,
+                    userHandle);
         }
     }
 }
diff --git a/src/com/android/settings/privatespace/PrivateSpaceSetupActivity.java b/src/com/android/settings/privatespace/PrivateSpaceSetupActivity.java
index 3a58e9e..a5628c8 100644
--- a/src/com/android/settings/privatespace/PrivateSpaceSetupActivity.java
+++ b/src/com/android/settings/privatespace/PrivateSpaceSetupActivity.java
@@ -31,6 +31,8 @@
 /** Activity class that helps in setting up of private space */
 public class PrivateSpaceSetupActivity extends FragmentActivity {
     public static final int SET_LOCK_ACTION = 1;
+    public static final int ACCOUNT_LOGIN_ACTION = 2;
+    public static final String EXTRA_ACTION_TYPE = "action_type";
     private NavHostFragment mNavHostFragment;
     @Override
     protected void onCreate(Bundle savedInstanceState) {
@@ -46,7 +48,13 @@
     @Override
     protected void onActivityResult(int requestCode, int resultCode, @Nullable Intent data) {
         if (requestCode == SET_LOCK_ACTION && resultCode == RESULT_OK) {
-            mNavHostFragment.getNavController().navigate(R.id.action_advance_to_success);
+            mNavHostFragment.getNavController().navigate(R.id.action_success_fragment);
+        } else if (requestCode == ACCOUNT_LOGIN_ACTION) {
+            if (resultCode == RESULT_OK) {
+                mNavHostFragment.getNavController().navigate(R.id.action_set_lock_fragment);
+            } else {
+                mNavHostFragment.getNavController().navigate(R.id.action_advance_login_error);
+            }
         }
         super.onActivityResult(requestCode, resultCode, data);
     }
diff --git a/src/com/android/settings/privatespace/SetupSuccessFragment.java b/src/com/android/settings/privatespace/SetupSuccessFragment.java
index a8ca3f1..b761da7 100644
--- a/src/com/android/settings/privatespace/SetupSuccessFragment.java
+++ b/src/com/android/settings/privatespace/SetupSuccessFragment.java
@@ -70,7 +70,7 @@
     private View.OnClickListener onClickNext() {
         return v -> {
             accessPrivateSpaceToast();
-            // TODO: Replace with the intent to launch PS/PS Launch Settings
+            // TODO(b/306228087): Replace with the intent to launch All Apps once it is working.
             Intent startMain = new Intent(Intent.ACTION_MAIN);
             startMain.addCategory(Intent.CATEGORY_HOME);
             startActivity(startMain);
diff --git a/tests/robotests/testutils/com/android/settings/testutils/FakeFeatureFactory.java b/tests/robotests/testutils/com/android/settings/testutils/FakeFeatureFactory.java
index 9156cae..5a5008c 100644
--- a/tests/robotests/testutils/com/android/settings/testutils/FakeFeatureFactory.java
+++ b/tests/robotests/testutils/com/android/settings/testutils/FakeFeatureFactory.java
@@ -46,6 +46,7 @@
 import com.android.settings.overlay.SupportFeatureProvider;
 import com.android.settings.overlay.SurveyFeatureProvider;
 import com.android.settings.panel.PanelFeatureProvider;
+import com.android.settings.privatespace.PrivateSpaceLoginFeatureProvider;
 import com.android.settings.search.SearchFeatureProvider;
 import com.android.settings.security.SecurityFeatureProvider;
 import com.android.settings.security.SecuritySettingsFeatureProvider;
@@ -99,6 +100,7 @@
     public StylusFeatureProvider mStylusFeatureProvider;
     public OnboardingFeatureProvider mOnboardingFeatureProvider;
     public FastPairFeatureProvider mFastPairFeatureProvider;
+    public PrivateSpaceLoginFeatureProvider mPrivateSpaceLoginFeatureProvider;
 
     /**
      * Call this in {@code @Before} method of the test class to use fake factory.
@@ -146,6 +148,7 @@
         mStylusFeatureProvider = mock(StylusFeatureProvider.class);
         mOnboardingFeatureProvider = mock(OnboardingFeatureProvider.class);
         mFastPairFeatureProvider = mock(FastPairFeatureProvider.class);
+        mPrivateSpaceLoginFeatureProvider =  mock(PrivateSpaceLoginFeatureProvider.class);
     }
 
     @Override
@@ -323,5 +326,10 @@
     public FastPairFeatureProvider getFastPairFeatureProvider() {
         return mFastPairFeatureProvider;
     }
+
+    @Override
+    public PrivateSpaceLoginFeatureProvider getPrivateSpaceLoginFeatureProvider() {
+        return mPrivateSpaceLoginFeatureProvider;
+    }
 }
 
diff --git a/tests/spa_unit/src/com/android/settings/testutils/FakeFeatureFactory.kt b/tests/spa_unit/src/com/android/settings/testutils/FakeFeatureFactory.kt
index 54299eb..9b098a7 100644
--- a/tests/spa_unit/src/com/android/settings/testutils/FakeFeatureFactory.kt
+++ b/tests/spa_unit/src/com/android/settings/testutils/FakeFeatureFactory.kt
@@ -41,6 +41,7 @@
 import com.android.settings.overlay.FeatureFactory
 import com.android.settings.overlay.SurveyFeatureProvider
 import com.android.settings.panel.PanelFeatureProvider
+import com.android.settings.privatespace.PrivateSpaceLoginFeatureProvider
 import com.android.settings.search.SearchFeatureProvider
 import com.android.settings.security.SecurityFeatureProvider
 import com.android.settings.security.SecuritySettingsFeatureProvider
@@ -143,4 +144,6 @@
         get() = TODO("Not yet implemented")
     override val fastPairFeatureProvider: FastPairFeatureProvider
         get() = TODO("Not yet implemented")
+    override val privateSpaceLoginFeatureProvider: PrivateSpaceLoginFeatureProvider
+        get() = TODO("Not yet implemented")
 }
diff --git a/tests/unit/src/com/android/settings/testutils/FakeFeatureFactory.java b/tests/unit/src/com/android/settings/testutils/FakeFeatureFactory.java
index b5062a0..bf2c84a 100644
--- a/tests/unit/src/com/android/settings/testutils/FakeFeatureFactory.java
+++ b/tests/unit/src/com/android/settings/testutils/FakeFeatureFactory.java
@@ -46,6 +46,7 @@
 import com.android.settings.overlay.SupportFeatureProvider;
 import com.android.settings.overlay.SurveyFeatureProvider;
 import com.android.settings.panel.PanelFeatureProvider;
+import com.android.settings.privatespace.PrivateSpaceLoginFeatureProvider;
 import com.android.settings.search.SearchFeatureProvider;
 import com.android.settings.security.SecurityFeatureProvider;
 import com.android.settings.security.SecuritySettingsFeatureProvider;
@@ -98,6 +99,7 @@
     public StylusFeatureProvider mStylusFeatureProvider;
     public OnboardingFeatureProvider mOnboardingFeatureProvider;
     public FastPairFeatureProvider mFastPairFeatureProvider;
+    public PrivateSpaceLoginFeatureProvider mPrivateSpaceLoginFeatureProvider;
 
     /**
      * Call this in {@code @Before} method of the test class to use fake factory.
@@ -145,6 +147,7 @@
         mStylusFeatureProvider = mock(StylusFeatureProvider.class);
         mOnboardingFeatureProvider = mock(OnboardingFeatureProvider.class);
         mFastPairFeatureProvider = mock(FastPairFeatureProvider.class);
+        mPrivateSpaceLoginFeatureProvider = mock(PrivateSpaceLoginFeatureProvider.class);
     }
 
     @Override
@@ -322,4 +325,9 @@
     public FastPairFeatureProvider getFastPairFeatureProvider() {
         return mFastPairFeatureProvider;
     }
+
+    @Override
+    public PrivateSpaceLoginFeatureProvider getPrivateSpaceLoginFeatureProvider() {
+        return mPrivateSpaceLoginFeatureProvider;
+    }
 }