Merge "Update DND quick affordance to work better with modes" into main
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/quickaffordance/DoNotDisturbQuickAffordanceConfigTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/quickaffordance/DoNotDisturbQuickAffordanceConfigTest.kt
index 6c3c7ef..fcf4662 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/quickaffordance/DoNotDisturbQuickAffordanceConfigTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/quickaffordance/DoNotDisturbQuickAffordanceConfigTest.kt
@@ -16,7 +16,10 @@
  */
 package com.android.systemui.keyguard.data.quickaffordance
 
+import android.app.Flags
 import android.net.Uri
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
 import android.provider.Settings
 import android.provider.Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS
 import android.provider.Settings.Global.ZEN_MODE_OFF
@@ -25,6 +28,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.settingslib.notification.modes.EnableZenModeDialog
+import com.android.settingslib.notification.modes.TestModeBuilder
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.Expandable
 import com.android.systemui.common.shared.model.ContentDescription
@@ -35,7 +39,11 @@
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.res.R
 import com.android.systemui.settings.UserTracker
+import com.android.systemui.shared.settings.data.repository.secureSettingsRepository
 import com.android.systemui.statusbar.policy.ZenModeController
+import com.android.systemui.statusbar.policy.data.repository.fakeDeviceProvisioningRepository
+import com.android.systemui.statusbar.policy.data.repository.fakeZenModeRepository
+import com.android.systemui.statusbar.policy.domain.interactor.zenModeInteractor
 import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.argumentCaptor
 import com.android.systemui.util.mockito.eq
@@ -43,6 +51,7 @@
 import com.android.systemui.util.mockito.whenever
 import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth.assertThat
+import java.time.Duration
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
@@ -66,8 +75,13 @@
     private val kosmos = testKosmos()
     private val testDispatcher = kosmos.testDispatcher
     private val testScope = kosmos.testScope
+
     private val settings = kosmos.fakeSettings
 
+    private val zenModeRepository = kosmos.fakeZenModeRepository
+    private val deviceProvisioningRepository = kosmos.fakeDeviceProvisioningRepository
+    private val secureSettingsRepository = kosmos.secureSettingsRepository
+
     @Mock private lateinit var zenModeController: ZenModeController
     @Mock private lateinit var userTracker: UserTracker
     @Mock private lateinit var conditionUri: Uri
@@ -85,17 +99,36 @@
             DoNotDisturbQuickAffordanceConfig(
                 context,
                 zenModeController,
+                kosmos.zenModeInteractor,
                 settings,
                 userTracker,
                 testDispatcher,
+                testScope.backgroundScope,
                 conditionUri,
                 enableZenModeDialog,
             )
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_MODES_UI)
     fun dndNotAvailable_pickerStateHidden() =
         testScope.runTest {
+            deviceProvisioningRepository.setDeviceProvisioned(false)
+            runCurrent()
+
+            val result = underTest.getPickerScreenState()
+            runCurrent()
+
+            assertEquals(
+                KeyguardQuickAffordanceConfig.PickerScreenState.UnavailableOnDevice,
+                result,
+            )
+        }
+
+    @Test
+    @DisableFlags(Flags.FLAG_MODES_UI)
+    fun controllerDndNotAvailable_pickerStateHidden() =
+        testScope.runTest {
             // given
             whenever(zenModeController.isZenAvailable).thenReturn(false)
 
@@ -105,13 +138,33 @@
             // then
             assertEquals(
                 KeyguardQuickAffordanceConfig.PickerScreenState.UnavailableOnDevice,
-                result
+                result,
             )
         }
 
     @Test
+    @EnableFlags(Flags.FLAG_MODES_UI)
     fun dndAvailable_pickerStateVisible() =
         testScope.runTest {
+            deviceProvisioningRepository.setDeviceProvisioned(true)
+            runCurrent()
+
+            val result = underTest.getPickerScreenState()
+            runCurrent()
+
+            assertThat(result)
+                .isInstanceOf(KeyguardQuickAffordanceConfig.PickerScreenState.Default::class.java)
+            val defaultPickerState =
+                result as KeyguardQuickAffordanceConfig.PickerScreenState.Default
+            assertThat(defaultPickerState.configureIntent).isNotNull()
+            assertThat(defaultPickerState.configureIntent?.action)
+                .isEqualTo(Settings.ACTION_ZEN_MODE_SETTINGS)
+        }
+
+    @Test
+    @DisableFlags(Flags.FLAG_MODES_UI)
+    fun controllerDndAvailable_pickerStateVisible() =
+        testScope.runTest {
             // given
             whenever(zenModeController.isZenAvailable).thenReturn(true)
 
@@ -129,7 +182,27 @@
         }
 
     @Test
-    fun onTriggered_dndModeIsNotZEN_MODE_OFF_setToZEN_MODE_OFF() =
+    @EnableFlags(Flags.FLAG_MODES_UI)
+    fun onTriggered_dndModeIsNotOff_setToOff() =
+        testScope.runTest {
+            val currentModes by collectLastValue(zenModeRepository.modes)
+
+            zenModeRepository.addMode(TestModeBuilder.MANUAL_DND_ACTIVE)
+            secureSettingsRepository.setInt(Settings.Secure.ZEN_DURATION, -2)
+            collectLastValue(underTest.lockScreenState)
+            runCurrent()
+
+            val result = underTest.onTriggered(null)
+            runCurrent()
+
+            val dndMode = currentModes!!.single()
+            assertThat(dndMode.isActive).isFalse()
+            assertEquals(KeyguardQuickAffordanceConfig.OnTriggeredResult.Handled, result)
+        }
+
+    @Test
+    @DisableFlags(Flags.FLAG_MODES_UI)
+    fun onTriggered_controllerDndModeIsNotZEN_MODE_OFF_setToZEN_MODE_OFF() =
         testScope.runTest {
             // given
             whenever(zenModeController.isZenAvailable).thenReturn(true)
@@ -140,11 +213,12 @@
 
             // when
             val result = underTest.onTriggered(null)
+
             verify(zenModeController)
                 .setZen(
                     spyZenMode.capture(),
                     spyConditionId.capture(),
-                    eq(DoNotDisturbQuickAffordanceConfig.TAG)
+                    eq(DoNotDisturbQuickAffordanceConfig.TAG),
                 )
 
             // then
@@ -154,7 +228,28 @@
         }
 
     @Test
-    fun onTriggered_dndModeIsZEN_MODE_OFF_settingFOREVER_setZenWithoutCondition() =
+    @EnableFlags(Flags.FLAG_MODES_UI)
+    fun onTriggered_dndModeIsOff_settingFOREVER_setZenWithoutCondition() =
+        testScope.runTest {
+            val currentModes by collectLastValue(zenModeRepository.modes)
+
+            zenModeRepository.addMode(TestModeBuilder.MANUAL_DND_INACTIVE)
+            secureSettingsRepository.setInt(Settings.Secure.ZEN_DURATION, ZEN_DURATION_FOREVER)
+            collectLastValue(underTest.lockScreenState)
+            runCurrent()
+
+            val result = underTest.onTriggered(null)
+            runCurrent()
+
+            val dndMode = currentModes!!.single()
+            assertThat(dndMode.isActive).isTrue()
+            assertThat(zenModeRepository.getModeActiveDuration(dndMode.id)).isNull()
+            assertEquals(KeyguardQuickAffordanceConfig.OnTriggeredResult.Handled, result)
+        }
+
+    @Test
+    @DisableFlags(Flags.FLAG_MODES_UI)
+    fun onTriggered_controllerDndModeIsZEN_MODE_OFF_settingFOREVER_setZenWithoutCondition() =
         testScope.runTest {
             // given
             whenever(zenModeController.isZenAvailable).thenReturn(true)
@@ -169,7 +264,7 @@
                 .setZen(
                     spyZenMode.capture(),
                     spyConditionId.capture(),
-                    eq(DoNotDisturbQuickAffordanceConfig.TAG)
+                    eq(DoNotDisturbQuickAffordanceConfig.TAG),
                 )
 
             // then
@@ -179,7 +274,27 @@
         }
 
     @Test
-    fun onTriggered_dndZEN_MODE_OFF_settingNotFOREVERorPROMPT_zenWithCondition() =
+    @EnableFlags(Flags.FLAG_MODES_UI)
+    fun onTriggered_dndModeIsOff_settingNotFOREVERorPROMPT_dndWithDuration() =
+        testScope.runTest {
+            val currentModes by collectLastValue(zenModeRepository.modes)
+            zenModeRepository.addMode(TestModeBuilder.MANUAL_DND_INACTIVE)
+            secureSettingsRepository.setInt(Settings.Secure.ZEN_DURATION, -900)
+            runCurrent()
+
+            val result = underTest.onTriggered(null)
+            runCurrent()
+
+            assertEquals(KeyguardQuickAffordanceConfig.OnTriggeredResult.Handled, result)
+            val dndMode = currentModes!!.single()
+            assertThat(dndMode.isActive).isTrue()
+            assertThat(zenModeRepository.getModeActiveDuration(dndMode.id))
+                .isEqualTo(Duration.ofMinutes(-900))
+        }
+
+    @Test
+    @DisableFlags(Flags.FLAG_MODES_UI)
+    fun onTriggered_controllerDndZEN_MODE_OFF_settingNotFOREVERorPROMPT_zenWithCondition() =
         testScope.runTest {
             // given
             whenever(zenModeController.isZenAvailable).thenReturn(true)
@@ -194,7 +309,7 @@
                 .setZen(
                     spyZenMode.capture(),
                     spyConditionId.capture(),
-                    eq(DoNotDisturbQuickAffordanceConfig.TAG)
+                    eq(DoNotDisturbQuickAffordanceConfig.TAG),
                 )
 
             // then
@@ -204,7 +319,28 @@
         }
 
     @Test
-    fun onTriggered_dndModeIsZEN_MODE_OFF_settingIsPROMPT_showDialog() =
+    @EnableFlags(Flags.FLAG_MODES_UI)
+    fun onTriggered_dndModeIsOff_settingIsPROMPT_showDialog() =
+        testScope.runTest {
+            val expandable: Expandable = mock()
+            zenModeRepository.addMode(TestModeBuilder.MANUAL_DND_INACTIVE)
+            secureSettingsRepository.setInt(Settings.Secure.ZEN_DURATION, ZEN_DURATION_PROMPT)
+            whenever(enableZenModeDialog.createDialog()).thenReturn(mock())
+            collectLastValue(underTest.lockScreenState)
+            runCurrent()
+
+            val result = underTest.onTriggered(expandable)
+
+            assertTrue(result is KeyguardQuickAffordanceConfig.OnTriggeredResult.ShowDialog)
+            assertEquals(
+                expandable,
+                (result as KeyguardQuickAffordanceConfig.OnTriggeredResult.ShowDialog).expandable,
+            )
+        }
+
+    @Test
+    @DisableFlags(Flags.FLAG_MODES_UI)
+    fun onTriggered_controllerDndModeIsZEN_MODE_OFF_settingIsPROMPT_showDialog() =
         testScope.runTest {
             // given
             val expandable: Expandable = mock()
@@ -222,13 +358,31 @@
             assertTrue(result is KeyguardQuickAffordanceConfig.OnTriggeredResult.ShowDialog)
             assertEquals(
                 expandable,
-                (result as KeyguardQuickAffordanceConfig.OnTriggeredResult.ShowDialog).expandable
+                (result as KeyguardQuickAffordanceConfig.OnTriggeredResult.ShowDialog).expandable,
             )
         }
 
     @Test
+    @EnableFlags(Flags.FLAG_MODES_UI)
     fun lockScreenState_dndAvailableStartsAsTrue_changeToFalse_StateIsHidden() =
         testScope.runTest {
+            deviceProvisioningRepository.setDeviceProvisioned(true)
+            val valueSnapshot = collectLastValue(underTest.lockScreenState)
+            val secondLastValue = valueSnapshot()
+            runCurrent()
+
+            deviceProvisioningRepository.setDeviceProvisioned(false)
+            runCurrent()
+            val lastValue = valueSnapshot()
+
+            assertTrue(secondLastValue is KeyguardQuickAffordanceConfig.LockScreenState.Visible)
+            assertTrue(lastValue is KeyguardQuickAffordanceConfig.LockScreenState.Hidden)
+        }
+
+    @Test
+    @DisableFlags(Flags.FLAG_MODES_UI)
+    fun lockScreenState_controllerDndAvailableStartsAsTrue_changeToFalse_StateIsHidden() =
+        testScope.runTest {
             // given
             whenever(zenModeController.isZenAvailable).thenReturn(true)
             val callbackCaptor: ArgumentCaptor<ZenModeController.Callback> = argumentCaptor()
@@ -246,7 +400,44 @@
         }
 
     @Test
-    fun lockScreenState_dndModeStartsAsZEN_MODE_OFF_changeToNotOFF_StateVisible() =
+    @EnableFlags(Flags.FLAG_MODES_UI)
+    fun lockScreenState_dndModeStartsAsOff_changeToOn_StateVisible() =
+        testScope.runTest {
+            val lockScreenState by collectLastValue(underTest.lockScreenState)
+
+            zenModeRepository.addMode(TestModeBuilder.MANUAL_DND_INACTIVE)
+            runCurrent()
+
+            assertThat(lockScreenState)
+                .isEqualTo(
+                    KeyguardQuickAffordanceConfig.LockScreenState.Visible(
+                        Icon.Resource(
+                            R.drawable.qs_dnd_icon_off,
+                            ContentDescription.Resource(R.string.dnd_is_off),
+                        ),
+                        ActivationState.Inactive,
+                    )
+                )
+
+            zenModeRepository.removeMode(TestModeBuilder.MANUAL_DND_INACTIVE.id)
+            zenModeRepository.addMode(TestModeBuilder.MANUAL_DND_ACTIVE)
+            runCurrent()
+
+            assertThat(lockScreenState)
+                .isEqualTo(
+                    KeyguardQuickAffordanceConfig.LockScreenState.Visible(
+                        Icon.Resource(
+                            R.drawable.qs_dnd_icon_on,
+                            ContentDescription.Resource(R.string.dnd_is_on),
+                        ),
+                        ActivationState.Active,
+                    )
+                )
+        }
+
+    @Test
+    @DisableFlags(Flags.FLAG_MODES_UI)
+    fun lockScreenState_controllerDndModeStartsAsZEN_MODE_OFF_changeToNotOFF_StateVisible() =
         testScope.runTest {
             // given
             whenever(zenModeController.isZenAvailable).thenReturn(true)
@@ -265,9 +456,9 @@
                 KeyguardQuickAffordanceConfig.LockScreenState.Visible(
                     Icon.Resource(
                         R.drawable.qs_dnd_icon_off,
-                        ContentDescription.Resource(R.string.dnd_is_off)
+                        ContentDescription.Resource(R.string.dnd_is_off),
                     ),
-                    ActivationState.Inactive
+                    ActivationState.Inactive,
                 ),
                 secondLastValue,
             )
@@ -275,9 +466,9 @@
                 KeyguardQuickAffordanceConfig.LockScreenState.Visible(
                     Icon.Resource(
                         R.drawable.qs_dnd_icon_on,
-                        ContentDescription.Resource(R.string.dnd_is_on)
+                        ContentDescription.Resource(R.string.dnd_is_on),
                     ),
-                    ActivationState.Active
+                    ActivationState.Active,
                 ),
                 lastValue,
             )
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorTest.kt
index fb32855..0f6dc07 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorTest.kt
@@ -17,7 +17,9 @@
 package com.android.systemui.statusbar.policy.domain.interactor
 
 import android.app.AutomaticZenRule
+import android.app.Flags
 import android.app.NotificationManager.Policy
+import android.platform.test.annotations.EnableFlags
 import android.provider.Settings
 import android.provider.Settings.Secure.ZEN_DURATION
 import android.provider.Settings.Secure.ZEN_DURATION_FOREVER
@@ -32,6 +34,7 @@
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.shared.settings.data.repository.secureSettingsRepository
+import com.android.systemui.statusbar.policy.data.repository.fakeDeviceProvisioningRepository
 import com.android.systemui.statusbar.policy.data.repository.fakeZenModeRepository
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
@@ -50,10 +53,31 @@
     private val testScope = kosmos.testScope
     private val zenModeRepository = kosmos.fakeZenModeRepository
     private val settingsRepository = kosmos.secureSettingsRepository
+    private val deviceProvisioningRepository = kosmos.fakeDeviceProvisioningRepository
 
     private val underTest = kosmos.zenModeInteractor
 
     @Test
+    fun isZenAvailable_off() =
+        testScope.runTest {
+            val isZenAvailable by collectLastValue(underTest.isZenAvailable)
+            deviceProvisioningRepository.setDeviceProvisioned(false)
+            runCurrent()
+
+            assertThat(isZenAvailable).isFalse()
+        }
+
+    @Test
+    fun isZenAvailable_on() =
+        testScope.runTest {
+            val isZenAvailable by collectLastValue(underTest.isZenAvailable)
+            deviceProvisioningRepository.setDeviceProvisioned(true)
+            runCurrent()
+
+            assertThat(isZenAvailable).isTrue()
+        }
+
+    @Test
     fun isZenModeEnabled_off() =
         testScope.runTest {
             val enabled by collectLastValue(underTest.isZenModeEnabled)
@@ -337,4 +361,22 @@
             runCurrent()
             assertThat(mainActiveMode).isNull()
         }
+
+    @Test
+    @EnableFlags(Flags.FLAG_MODES_UI)
+    fun dndMode_flows() =
+        testScope.runTest {
+            val dndMode by collectLastValue(underTest.dndMode)
+
+            zenModeRepository.addMode(TestModeBuilder.MANUAL_DND_INACTIVE)
+            runCurrent()
+
+            assertThat(dndMode!!.isActive).isFalse()
+
+            zenModeRepository.removeMode(TestModeBuilder.MANUAL_DND_INACTIVE.id)
+            zenModeRepository.addMode(TestModeBuilder.MANUAL_DND_ACTIVE)
+            runCurrent()
+
+            assertThat(dndMode!!.isActive).isTrue()
+        }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/data/quickaffordance/DoNotDisturbQuickAffordanceConfig.kt b/packages/SystemUI/src/com/android/systemui/keyguard/data/quickaffordance/DoNotDisturbQuickAffordanceConfig.kt
index 406b9f6..be87334 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/data/quickaffordance/DoNotDisturbQuickAffordanceConfig.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/data/quickaffordance/DoNotDisturbQuickAffordanceConfig.kt
@@ -25,7 +25,9 @@
 import android.provider.Settings.Secure.ZEN_DURATION_FOREVER
 import android.provider.Settings.Secure.ZEN_DURATION_PROMPT
 import android.service.notification.ZenModeConfig
+import android.util.Log
 import com.android.settingslib.notification.modes.EnableZenModeDialog
+import com.android.settingslib.notification.modes.ZenMode
 import com.android.settingslib.notification.modes.ZenModeDialogMetricsLogger
 import com.android.systemui.animation.Expandable
 import com.android.systemui.common.coroutine.ChannelExt.trySendWithFailureLogging
@@ -35,30 +37,38 @@
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.keyguard.shared.quickaffordance.ActivationState
+import com.android.systemui.modes.shared.ModesUi
 import com.android.systemui.res.R
 import com.android.systemui.settings.UserTracker
 import com.android.systemui.statusbar.policy.ZenModeController
+import com.android.systemui.statusbar.policy.domain.interactor.ZenModeInteractor
 import com.android.systemui.util.settings.SecureSettings
 import com.android.systemui.util.settings.SettingsProxyExt.observerFlow
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.onStart
+import kotlinx.coroutines.flow.stateIn
 
 @SysUISingleton
 class DoNotDisturbQuickAffordanceConfig
 constructor(
     private val context: Context,
     private val controller: ZenModeController,
+    private val interactor: ZenModeInteractor,
     private val secureSettings: SecureSettings,
     private val userTracker: UserTracker,
     @Background private val backgroundDispatcher: CoroutineDispatcher,
+    @Background private val backgroundScope: CoroutineScope,
     private val testConditionId: Uri?,
     testDialog: EnableZenModeDialog?,
 ) : KeyguardQuickAffordanceConfig {
@@ -67,15 +77,45 @@
     constructor(
         context: Context,
         controller: ZenModeController,
+        interactor: ZenModeInteractor,
         secureSettings: SecureSettings,
         userTracker: UserTracker,
         @Background backgroundDispatcher: CoroutineDispatcher,
-    ) : this(context, controller, secureSettings, userTracker, backgroundDispatcher, null, null)
+        @Background backgroundScope: CoroutineScope,
+    ) : this(
+        context,
+        controller,
+        interactor,
+        secureSettings,
+        userTracker,
+        backgroundDispatcher,
+        backgroundScope,
+        null,
+        null,
+    )
 
-    private var dndMode: Int = 0
-    private var isAvailable = false
+    private var zenMode: Int = 0
+    private var oldIsAvailable = false
     private var settingsValue: Int = 0
 
+    private val dndMode: StateFlow<ZenMode?> by lazy {
+        ModesUi.assertInNewMode()
+        interactor.dndMode.stateIn(
+            scope = backgroundScope,
+            started = SharingStarted.Eagerly,
+            initialValue = null,
+        )
+    }
+
+    private val isAvailable: StateFlow<Boolean> by lazy {
+        ModesUi.assertInNewMode()
+        interactor.isZenAvailable.stateIn(
+            scope = backgroundScope,
+            started = SharingStarted.Eagerly,
+            initialValue = false,
+        )
+    }
+
     private val conditionUri: Uri
         get() =
             testConditionId
@@ -104,42 +144,68 @@
     override val pickerIconResourceId: Int = R.drawable.ic_do_not_disturb
 
     override val lockScreenState: Flow<KeyguardQuickAffordanceConfig.LockScreenState> =
-        combine(
-            conflatedCallbackFlow {
-                val callback =
-                    object : ZenModeController.Callback {
-                        override fun onZenChanged(zen: Int) {
-                            dndMode = zen
-                            trySendWithFailureLogging(updateState(), TAG)
+        if (ModesUi.isEnabled) {
+            combine(isAvailable, dndMode) { isAvailable, dndMode ->
+                if (!isAvailable) {
+                    KeyguardQuickAffordanceConfig.LockScreenState.Hidden
+                } else if (dndMode?.isActive == true) {
+                    KeyguardQuickAffordanceConfig.LockScreenState.Visible(
+                        Icon.Resource(
+                            R.drawable.qs_dnd_icon_on,
+                            ContentDescription.Resource(R.string.dnd_is_on),
+                        ),
+                        ActivationState.Active,
+                    )
+                } else {
+                    KeyguardQuickAffordanceConfig.LockScreenState.Visible(
+                        Icon.Resource(
+                            R.drawable.qs_dnd_icon_off,
+                            ContentDescription.Resource(R.string.dnd_is_off),
+                        ),
+                        ActivationState.Inactive,
+                    )
+                }
+            }
+        } else {
+            combine(
+                conflatedCallbackFlow {
+                    val callback =
+                        object : ZenModeController.Callback {
+                            override fun onZenChanged(zen: Int) {
+                                zenMode = zen
+                                trySendWithFailureLogging(updateState(), TAG)
+                            }
+
+                            override fun onZenAvailableChanged(available: Boolean) {
+                                oldIsAvailable = available
+                                trySendWithFailureLogging(updateState(), TAG)
+                            }
                         }
 
-                        override fun onZenAvailableChanged(available: Boolean) {
-                            isAvailable = available
-                            trySendWithFailureLogging(updateState(), TAG)
-                        }
-                    }
+                    zenMode = controller.zen
+                    oldIsAvailable = controller.isZenAvailable
+                    trySendWithFailureLogging(updateState(), TAG)
 
-                dndMode = controller.zen
-                isAvailable = controller.isZenAvailable
-                trySendWithFailureLogging(updateState(), TAG)
+                    controller.addCallback(callback)
 
-                controller.addCallback(callback)
-
-                awaitClose { controller.removeCallback(callback) }
-            },
-            secureSettings
-                .observerFlow(userTracker.userId, Settings.Secure.ZEN_DURATION)
-                .onStart { emit(Unit) }
-                .map { secureSettings.getInt(Settings.Secure.ZEN_DURATION, ZEN_MODE_OFF) }
-                .flowOn(backgroundDispatcher)
-                .distinctUntilChanged()
-                .onEach { settingsValue = it }
-        ) { callbackFlowValue, _ ->
-            callbackFlowValue
+                    awaitClose { controller.removeCallback(callback) }
+                },
+                secureSettings
+                    .observerFlow(userTracker.userId, Settings.Secure.ZEN_DURATION)
+                    .onStart { emit(Unit) }
+                    .map { secureSettings.getInt(Settings.Secure.ZEN_DURATION, ZEN_MODE_OFF) }
+                    .flowOn(backgroundDispatcher)
+                    .distinctUntilChanged()
+                    .onEach { settingsValue = it },
+            ) { callbackFlowValue, _ ->
+                callbackFlowValue
+            }
         }
 
     override suspend fun getPickerScreenState(): KeyguardQuickAffordanceConfig.PickerScreenState {
-        return if (controller.isZenAvailable) {
+        val isZenAvailable = if (ModesUi.isEnabled) isAvailable.value else controller.isZenAvailable
+
+        return if (isZenAvailable) {
             KeyguardQuickAffordanceConfig.PickerScreenState.Default(
                 configureIntent = Intent(Settings.ACTION_ZEN_MODE_SETTINGS)
             )
@@ -151,32 +217,63 @@
     override fun onTriggered(
         expandable: Expandable?
     ): KeyguardQuickAffordanceConfig.OnTriggeredResult {
-        return when {
-            !isAvailable -> KeyguardQuickAffordanceConfig.OnTriggeredResult.Handled
-            dndMode != ZEN_MODE_OFF -> {
-                controller.setZen(ZEN_MODE_OFF, null, TAG)
+        return if (ModesUi.isEnabled) {
+            if (!isAvailable.value) {
                 KeyguardQuickAffordanceConfig.OnTriggeredResult.Handled
+            } else {
+                val dnd = dndMode.value
+                if (dnd == null) {
+                    Log.wtf(TAG, "Triggered DND but it's null!?")
+                    return KeyguardQuickAffordanceConfig.OnTriggeredResult.Handled
+                }
+                if (dnd.isActive) {
+                    interactor.deactivateMode(dnd)
+                    return KeyguardQuickAffordanceConfig.OnTriggeredResult.Handled
+                } else {
+                    if (interactor.shouldAskForZenDuration(dnd)) {
+                        // NOTE: The dialog handles turning on the mode itself.
+                        return KeyguardQuickAffordanceConfig.OnTriggeredResult.ShowDialog(
+                            dialog.createDialog(),
+                            expandable,
+                        )
+                    } else {
+                        interactor.activateMode(dnd)
+                        return KeyguardQuickAffordanceConfig.OnTriggeredResult.Handled
+                    }
+                }
             }
-            settingsValue == ZEN_DURATION_PROMPT ->
-                KeyguardQuickAffordanceConfig.OnTriggeredResult.ShowDialog(
-                    dialog.createDialog(),
-                    expandable
-                )
-            settingsValue == ZEN_DURATION_FOREVER -> {
-                controller.setZen(ZEN_MODE_IMPORTANT_INTERRUPTIONS, null, TAG)
-                KeyguardQuickAffordanceConfig.OnTriggeredResult.Handled
-            }
-            else -> {
-                controller.setZen(ZEN_MODE_IMPORTANT_INTERRUPTIONS, conditionUri, TAG)
-                KeyguardQuickAffordanceConfig.OnTriggeredResult.Handled
+        } else {
+            when {
+                !oldIsAvailable -> KeyguardQuickAffordanceConfig.OnTriggeredResult.Handled
+                zenMode != ZEN_MODE_OFF -> {
+                    controller.setZen(ZEN_MODE_OFF, null, TAG)
+                    KeyguardQuickAffordanceConfig.OnTriggeredResult.Handled
+                }
+
+                settingsValue == ZEN_DURATION_PROMPT ->
+                    KeyguardQuickAffordanceConfig.OnTriggeredResult.ShowDialog(
+                        dialog.createDialog(),
+                        expandable,
+                    )
+
+                settingsValue == ZEN_DURATION_FOREVER -> {
+                    controller.setZen(ZEN_MODE_IMPORTANT_INTERRUPTIONS, null, TAG)
+                    KeyguardQuickAffordanceConfig.OnTriggeredResult.Handled
+                }
+
+                else -> {
+                    controller.setZen(ZEN_MODE_IMPORTANT_INTERRUPTIONS, conditionUri, TAG)
+                    KeyguardQuickAffordanceConfig.OnTriggeredResult.Handled
+                }
             }
         }
     }
 
     private fun updateState(): KeyguardQuickAffordanceConfig.LockScreenState {
-        return if (!isAvailable) {
+        ModesUi.assertInLegacyMode()
+        return if (!oldIsAvailable) {
             KeyguardQuickAffordanceConfig.LockScreenState.Hidden
-        } else if (dndMode == ZEN_MODE_OFF) {
+        } else if (zenMode == ZEN_MODE_OFF) {
             KeyguardQuickAffordanceConfig.LockScreenState.Visible(
                 Icon.Resource(
                     R.drawable.qs_dnd_icon_off,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt
index dbeaa59..ba45942 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt
@@ -27,7 +27,10 @@
 import com.android.settingslib.notification.modes.ZenIconLoader
 import com.android.settingslib.notification.modes.ZenMode
 import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.modes.shared.ModesUi
 import com.android.systemui.shared.notifications.data.repository.NotificationSettingsRepository
+import com.android.systemui.statusbar.policy.data.repository.DeviceProvisioningRepository
+import com.android.systemui.statusbar.policy.data.repository.UserSetupRepository
 import com.android.systemui.statusbar.policy.domain.model.ActiveZenModes
 import com.android.systemui.statusbar.policy.domain.model.ZenModeInfo
 import java.time.Duration
@@ -51,7 +54,17 @@
     private val notificationSettingsRepository: NotificationSettingsRepository,
     @Background private val bgDispatcher: CoroutineDispatcher,
     private val iconLoader: ZenIconLoader,
+    private val deviceProvisioningRepository: DeviceProvisioningRepository,
+    private val userSetupRepository: UserSetupRepository,
 ) {
+    val isZenAvailable: Flow<Boolean> =
+        combine(
+            deviceProvisioningRepository.isDeviceProvisioned,
+            userSetupRepository.isUserSetUp,
+        ) { isDeviceProvisioned, isUserSetUp ->
+            isDeviceProvisioned && isUserSetUp
+        }
+
     val isZenModeEnabled: Flow<Boolean> =
         zenModeRepository.globalZenMode
             .map {
@@ -80,6 +93,18 @@
 
     val modes: Flow<List<ZenMode>> = zenModeRepository.modes
 
+    /**
+     * Returns the special "manual DND" mode.
+     *
+     * This is only meant as a temporary solution for "legacy" UI pieces that handle DND
+     * specifically; any new or migrated features should use modes more generally, through [modes]
+     * or [activeModes].
+     */
+    val dndMode: Flow<ZenMode?> by lazy {
+        ModesUi.assertInNewMode()
+        zenModeRepository.modes.map { modes -> modes.singleOrNull { it.isManualDnd } }
+    }
+
     /** Flow returning the currently active mode(s), if any. */
     val activeModes: Flow<ActiveZenModes> =
         modes
@@ -113,10 +138,11 @@
                         Log.e(
                             TAG,
                             "Interactor cannot handle showing the zen duration prompt. " +
-                                "Please use EnableZenModeDialog when this setting is active."
+                                "Please use EnableZenModeDialog when this setting is active.",
                         )
                         null
                     }
+
                     ZEN_DURATION_FOREVER -> null
                     else -> Duration.ofMinutes(zenDuration.toLong())
                 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorKosmos.kt
index 61b53c9..99cd830 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorKosmos.kt
@@ -22,6 +22,8 @@
 import com.android.systemui.kosmos.Kosmos.Fixture
 import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.shared.notifications.data.repository.notificationSettingsRepository
+import com.android.systemui.statusbar.policy.data.repository.deviceProvisioningRepository
+import com.android.systemui.statusbar.policy.data.repository.userSetupRepository
 import com.android.systemui.statusbar.policy.data.repository.zenModeRepository
 
 val Kosmos.zenModeInteractor by Fixture {
@@ -31,5 +33,7 @@
         notificationSettingsRepository = notificationSettingsRepository,
         bgDispatcher = testDispatcher,
         iconLoader = zenIconLoader,
+        deviceProvisioningRepository = deviceProvisioningRepository,
+        userSetupRepository = userSetupRepository,
     )
 }