Merge "Preselect widget that is long pressed when entering edit mode" into main
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt
index 761e74e..ee9a52c 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt
@@ -130,12 +130,11 @@
     val gridState = rememberLazyGridState()
     val contentListState = rememberContentListState(widgetConfigurator, communalContent, viewModel)
     val reorderingWidgets by viewModel.reorderingWidgets.collectAsState()
-    val selectedIndex = viewModel.selectedIndex.collectAsState()
+    val selectedKey = viewModel.selectedKey.collectAsState()
     val removeButtonEnabled by remember {
-        derivedStateOf { selectedIndex.value != null || reorderingWidgets }
+        derivedStateOf { selectedKey.value != null || reorderingWidgets }
     }
-    val (isButtonToEditWidgetsShowing, setIsButtonToEditWidgetsShowing) =
-        remember { mutableStateOf(false) }
+    var isButtonToEditWidgetsShowing by remember { mutableStateOf(false) }
 
     val contentPadding = gridContentPadding(viewModel.isEditMode, toolbarSize)
     val contentOffset = beforeContentPadding(contentPadding).toOffset()
@@ -150,22 +149,30 @@
                     if (!viewModel.isEditMode) return@pointerInput
                     observeTapsWithoutConsuming { offset ->
                         val adjustedOffset = offset - contentOffset
-                        val index =
-                            gridState.layoutInfo.visibleItemsInfo
-                                .firstItemAtOffset(adjustedOffset)
-                                ?.index
-                        val newIndex =
-                            if (index?.let(contentListState::isItemEditable) == true) {
-                                index
-                            } else {
-                                null
-                            }
-                        viewModel.setSelectedIndex(newIndex)
+                        val index = firstIndexAtOffset(gridState, adjustedOffset)
+                        val key = index?.let { keyAtIndexIfEditable(contentListState.list, index) }
+                        viewModel.setSelectedKey(key)
                     }
                 }
                 .thenIf(!viewModel.isEditMode) {
-                    Modifier.pointerInput(Unit) {
-                        detectLongPressGesture { offset -> setIsButtonToEditWidgetsShowing(true) }
+                    Modifier.pointerInput(
+                        gridState,
+                        contentOffset,
+                        communalContent,
+                        gridCoordinates
+                    ) {
+                        detectLongPressGesture { offset ->
+                            isButtonToEditWidgetsShowing = true
+
+                            // Deduct both grid offset relative to its container and content offset.
+                            val adjustedOffset =
+                                gridCoordinates?.let {
+                                    offset - it.positionInWindow() - contentOffset
+                                }
+                            val index = adjustedOffset?.let { firstIndexAtOffset(gridState, it) }
+                            val key = index?.let { keyAtIndexIfEditable(communalContent, index) }
+                            viewModel.setSelectedKey(key)
+                        }
                     }
                 },
     ) {
@@ -186,7 +193,7 @@
             onOpenWidgetPicker = onOpenWidgetPicker,
             gridState = gridState,
             contentListState = contentListState,
-            selectedIndex = selectedIndex,
+            selectedKey = selectedKey,
             widgetConfigurator = widgetConfigurator,
         )
 
@@ -198,10 +205,14 @@
                 onEditDone = onEditDone,
                 onOpenWidgetPicker = onOpenWidgetPicker,
                 onRemoveClicked = {
-                    selectedIndex.value?.let { index ->
-                        contentListState.onRemove(index)
+                    val index =
+                        selectedKey.value?.let { key ->
+                            contentListState.list.indexOfFirst { it.key == key }
+                        }
+                    index?.let {
+                        contentListState.onRemove(it)
                         contentListState.onSaveList()
-                        viewModel.setSelectedIndex(null)
+                        viewModel.setSelectedKey(null)
                     }
                 },
                 removeEnabled = removeButtonEnabled
@@ -219,10 +230,10 @@
         if (isButtonToEditWidgetsShowing) {
             ButtonToEditWidgets(
                 onClick = {
-                    setIsButtonToEditWidgetsShowing(false)
-                    viewModel.onOpenWidgetEditor()
+                    isButtonToEditWidgetsShowing = false
+                    viewModel.onOpenWidgetEditor(selectedKey.value)
                 },
-                onHide = { setIsButtonToEditWidgetsShowing(false) },
+                onHide = { isButtonToEditWidgetsShowing = false },
             )
         }
 
@@ -244,7 +255,7 @@
     communalContent: List<CommunalContentModel>,
     viewModel: BaseCommunalViewModel,
     contentPadding: PaddingValues,
-    selectedIndex: State<Int?>,
+    selectedKey: State<String?>,
     contentOffset: Offset,
     gridState: LazyGridState,
     contentListState: ContentListState,
@@ -253,7 +264,8 @@
     onOpenWidgetPicker: (() -> Unit)? = null,
     widgetConfigurator: WidgetConfigurator?,
 ) {
-    var gridModifier = Modifier.align(Alignment.CenterStart)
+    var gridModifier =
+        Modifier.align(Alignment.CenterStart).onGloballyPositioned { setGridCoordinates(it) }
     var list = communalContent
     var dragDropState: GridDragDropState? = null
     if (viewModel.isEditMode && viewModel is CommunalEditModeViewModel) {
@@ -266,10 +278,7 @@
                 updateDragPositionForRemove = updateDragPositionForRemove
             )
         gridModifier =
-            gridModifier
-                .fillMaxSize()
-                .dragContainer(dragDropState, contentOffset, viewModel)
-                .onGloballyPositioned { setGridCoordinates(it) }
+            gridModifier.fillMaxSize().dragContainer(dragDropState, contentOffset, viewModel)
         // for widgets dropped from other activities
         val dragAndDropTargetState =
             rememberDragAndDropTargetState(
@@ -307,7 +316,8 @@
                     list[index].size.dp().value,
                 )
             if (viewModel.isEditMode && dragDropState != null) {
-                val selected by remember(index) { derivedStateOf { index == selectedIndex.value } }
+                val selected by
+                    remember(index) { derivedStateOf { list[index].key == selectedKey.value } }
                 DraggableItem(
                     dragDropState = dragDropState,
                     selected = selected,
@@ -832,6 +842,13 @@
     }
 }
 
+private fun firstIndexAtOffset(gridState: LazyGridState, offset: Offset): Int? =
+    gridState.layoutInfo.visibleItemsInfo.firstItemAtOffset(offset)?.index
+
+/** Returns the key of item if it's editable at the given index. Only widget is editable. */
+private fun keyAtIndexIfEditable(list: List<CommunalContentModel>, index: Int): String? =
+    if (index in list.indices && list[index].isWidget()) list[index].key else null
+
 data class ContentPaddingInPx(val start: Float, val top: Float) {
     fun toOffset(): Offset = Offset(start, top)
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalInteractorTest.kt
index a083e7c..b1224ffa 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalInteractorTest.kt
@@ -633,6 +633,14 @@
             verify(editWidgetsActivityStarter).startActivity()
         }
 
+    @Test
+    fun showWidgetEditor_withPreselectedKey_startsActivity() =
+        testScope.runTest {
+            val widgetKey = CommunalContentModel.KEY.widget(123)
+            underTest.showWidgetEditor(preselectedKey = widgetKey)
+            verify(editWidgetsActivityStarter).startActivity(widgetKey)
+        }
+
     private fun smartspaceTimer(id: String, timestamp: Long = 0L): SmartspaceTarget {
         val timer = mock(SmartspaceTarget::class.java)
         whenever(timer.smartspaceTargetId).thenReturn(id)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt
index a2dec5f..273d1cd 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt
@@ -129,6 +129,19 @@
         }
 
     @Test
+    fun selectedKey_onReorderWidgets_isCleared() =
+        testScope.runTest {
+            val selectedKey by collectLastValue(underTest.selectedKey)
+
+            val key = CommunalContentModel.KEY.widget(123)
+            underTest.setSelectedKey(key)
+            assertThat(selectedKey).isEqualTo(key)
+
+            underTest.onReorderWidgetStart()
+            assertThat(selectedKey).isNull()
+        }
+
+    @Test
     fun reorderWidget_uiEventLogging_start() {
         underTest.onReorderWidgetStart()
         verify(uiEventLogger).log(CommunalUiEvent.COMMUNAL_HUB_REORDER_WIDGET_START)
diff --git a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
index 28adb77..89e31f3 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
@@ -176,8 +176,8 @@
     }
 
     /** Show the widget editor Activity. */
-    fun showWidgetEditor() {
-        editWidgetsActivityStarter.startActivity()
+    fun showWidgetEditor(preselectedKey: String? = null) {
+        editWidgetsActivityStarter.startActivity(preselectedKey)
     }
 
     /** Dismiss the CTA tile from the hub in view mode. */
diff --git a/packages/SystemUI/src/com/android/systemui/communal/domain/model/CommunalContentModel.kt b/packages/SystemUI/src/com/android/systemui/communal/domain/model/CommunalContentModel.kt
index acd6cb0..ae019a1 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/domain/model/CommunalContentModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/domain/model/CommunalContentModel.kt
@@ -128,4 +128,6 @@
             }
         }
     }
+
+    fun isWidget() = this is Widget
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/BaseCommunalViewModel.kt b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/BaseCommunalViewModel.kt
index 1e64d3f..91df828 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/BaseCommunalViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/BaseCommunalViewModel.kt
@@ -40,11 +40,11 @@
     /** Whether widgets are currently being re-ordered. */
     open val reorderingWidgets: StateFlow<Boolean> = MutableStateFlow(false)
 
-    private val _selectedIndex: MutableStateFlow<Int?> = MutableStateFlow(null)
+    private val _selectedKey: MutableStateFlow<String?> = MutableStateFlow(null)
 
-    /** The index of the currently selected item, or null if no item selected. */
-    val selectedIndex: StateFlow<Int?>
-        get() = _selectedIndex
+    /** The key of the currently selected item, or null if no item selected. */
+    val selectedKey: StateFlow<String?>
+        get() = _selectedKey
 
     fun onSceneChanged(scene: CommunalSceneKey) {
         communalInteractor.onSceneChanged(scene)
@@ -94,8 +94,8 @@
      */
     open fun onReorderWidgets(widgetIdToPriorityMap: Map<Int, Int>) {}
 
-    /** Called as the UI requests opening the widget editor. */
-    open fun onOpenWidgetEditor() {}
+    /** Called as the UI requests opening the widget editor with an optional preselected widget. */
+    open fun onOpenWidgetEditor(preselectedKey: String? = null) {}
 
     /** Called as the UI requests to dismiss the CTA tile. */
     open fun onDismissCtaTile() {}
@@ -109,8 +109,8 @@
     /** Called as the user cancels dragging a widget to reorder. */
     open fun onReorderWidgetCancel() {}
 
-    /** Set the index of the currently selected item */
-    fun setSelectedIndex(index: Int?) {
-        _selectedIndex.value = index
+    /** Set the key of the currently selected item */
+    fun setSelectedKey(key: String?) {
+        _selectedKey.value = key
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt
index 4b98f1a..ebcfb8b 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt
@@ -29,7 +29,6 @@
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.onEach
 
 /** The view model for communal hub in edit mode. */
 @SysUISingleton
@@ -44,10 +43,9 @@
 
     // Only widgets are editable. The CTA tile comes last in the list and remains visible.
     override val communalContent: Flow<List<CommunalContentModel>> =
-        communalInteractor.widgetContent
-            // Clear the selected index when the list is updated.
-            .onEach { setSelectedIndex(null) }
-            .map { widgets -> widgets + listOf(CommunalContentModel.CtaTileInEditMode()) }
+        communalInteractor.widgetContent.map { widgets ->
+            widgets + listOf(CommunalContentModel.CtaTileInEditMode())
+        }
 
     private val _reorderingWidgets = MutableStateFlow(false)
 
@@ -61,7 +59,7 @@
 
     override fun onReorderWidgetStart() {
         // Clear selection status
-        setSelectedIndex(null)
+        setSelectedKey(null)
         _reorderingWidgets.value = true
         uiEventLogger.log(CommunalUiEvent.COMMUNAL_HUB_REORDER_WIDGET_START)
     }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalViewModel.kt b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalViewModel.kt
index a909383..d7a3705 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalViewModel.kt
@@ -81,7 +81,8 @@
         }
     }
 
-    override fun onOpenWidgetEditor() = communalInteractor.showWidgetEditor()
+    override fun onOpenWidgetEditor(preselectedKey: String?) =
+        communalInteractor.showWidgetEditor(preselectedKey)
 
     override fun onDismissCtaTile() {
         scope.launch {
diff --git a/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt b/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt
index a257543..ad1327e 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt
@@ -47,6 +47,7 @@
         private const val EXTRA_FILTER_STRATEGY = "filter_strategy"
         private const val FILTER_STRATEGY_GLANCEABLE_HUB = 1
         private const val TAG = "EditWidgetsActivity"
+        const val EXTRA_PRESELECTED_KEY = "preselected_key"
     }
 
     private val widgetConfigurator by lazy { widgetConfiguratorFactory.create(this) }
@@ -92,6 +93,9 @@
         windowInsetsController?.hide(WindowInsets.Type.systemBars())
         window.setDecorFitsSystemWindows(false)
 
+        val preselectedKey = intent.getStringExtra(EXTRA_PRESELECTED_KEY)
+        communalViewModel.setSelectedKey(preselectedKey)
+
         setCommunalEditWidgetActivityContent(
             activity = this,
             viewModel = communalViewModel,
diff --git a/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivityStarter.kt b/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivityStarter.kt
index 55acad0..d1843af 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivityStarter.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivityStarter.kt
@@ -18,12 +18,13 @@
 
 import android.content.Context
 import android.content.Intent
+import com.android.systemui.communal.widgets.EditWidgetsActivity.Companion.EXTRA_PRESELECTED_KEY
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.plugins.ActivityStarter
 import javax.inject.Inject
 
 interface EditWidgetsActivityStarter {
-    fun startActivity()
+    fun startActivity(preselectedKey: String? = null)
 }
 
 class EditWidgetsActivityStarterImpl
@@ -33,10 +34,11 @@
     private val activityStarter: ActivityStarter,
 ) : EditWidgetsActivityStarter {
 
-    override fun startActivity() {
+    override fun startActivity(preselectedKey: String?) {
         activityStarter.startActivityDismissingKeyguard(
             Intent(applicationContext, EditWidgetsActivity::class.java)
-                .addFlags(Intent.FLAG_ACTIVITY_NEW_TASK or Intent.FLAG_ACTIVITY_CLEAR_TASK),
+                .addFlags(Intent.FLAG_ACTIVITY_NEW_TASK or Intent.FLAG_ACTIVITY_CLEAR_TASK)
+                .apply { preselectedKey?.let { putExtra(EXTRA_PRESELECTED_KEY, preselectedKey) } },
             /* onlyProvisioned = */ true,
             /* dismissShade = */ true,
         )