Added a toolbar in glanceable hub for editing

- Added a toolbar above the grid with three buttons for add/remove/done.
- Drag and drop a widget into the boundaries of the remove button
  will remove the item.
- TODO: clicking on Done button should go back to glanceable hub.

Bug: b/315154364
Test: manually
Flag: ACONFIG com.android.systemui.communal_hub DEVELOPMENT
Change-Id: Ib92026e9f43e0da36a50099bcb41f73c4cae7de9
diff --git a/packages/SystemUI/compose/facade/disabled/src/com/android/systemui/compose/ComposeFacade.kt b/packages/SystemUI/compose/facade/disabled/src/com/android/systemui/compose/ComposeFacade.kt
index 914e5f2..fd04b5ee0 100644
--- a/packages/SystemUI/compose/facade/disabled/src/com/android/systemui/compose/ComposeFacade.kt
+++ b/packages/SystemUI/compose/facade/disabled/src/com/android/systemui/compose/ComposeFacade.kt
@@ -51,6 +51,7 @@
         activity: ComponentActivity,
         viewModel: BaseCommunalViewModel,
         onOpenWidgetPicker: () -> Unit,
+        onEditDone: () -> Unit,
     ) {
         throwComposeUnavailableError()
     }
diff --git a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/compose/ComposeFacade.kt b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/compose/ComposeFacade.kt
index 59bd95b..5055ee1 100644
--- a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/compose/ComposeFacade.kt
+++ b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/compose/ComposeFacade.kt
@@ -66,12 +66,14 @@
         activity: ComponentActivity,
         viewModel: BaseCommunalViewModel,
         onOpenWidgetPicker: () -> Unit,
+        onEditDone: () -> Unit,
     ) {
         activity.setContent {
             PlatformTheme {
                 CommunalHub(
                     viewModel = viewModel,
                     onOpenWidgetPicker = onOpenWidgetPicker,
+                    onEditDone = onEditDone,
                 )
             }
         }
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 e8ecd3a..2a9cf0f 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
@@ -25,10 +25,12 @@
 import androidx.compose.foundation.layout.Arrangement
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.PaddingValues
+import androidx.compose.foundation.layout.Row
 import androidx.compose.foundation.layout.Spacer
-import androidx.compose.foundation.layout.fillMaxHeight
 import androidx.compose.foundation.layout.fillMaxSize
+import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.height
+import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.layout.width
 import androidx.compose.foundation.lazy.grid.GridCells
 import androidx.compose.foundation.lazy.grid.GridItemSpan
@@ -36,23 +38,41 @@
 import androidx.compose.foundation.lazy.grid.rememberLazyGridState
 import androidx.compose.material.icons.Icons
 import androidx.compose.material.icons.filled.Add
-import androidx.compose.material.icons.filled.Close
 import androidx.compose.material.icons.filled.Edit
+import androidx.compose.material.icons.outlined.Delete
+import androidx.compose.material3.Button
+import androidx.compose.material3.ButtonColors
+import androidx.compose.material3.ButtonDefaults
 import androidx.compose.material3.Card
 import androidx.compose.material3.CardDefaults
 import androidx.compose.material3.Icon
 import androidx.compose.material3.IconButton
+import androidx.compose.material3.OutlinedButton
+import androidx.compose.material3.Text
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.collectAsState
 import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
+import androidx.compose.runtime.setValue
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.input.pointer.pointerInput
+import androidx.compose.ui.layout.LayoutCoordinates
+import androidx.compose.ui.layout.onGloballyPositioned
+import androidx.compose.ui.layout.onSizeChanged
+import androidx.compose.ui.layout.positionInWindow
+import androidx.compose.ui.platform.LocalConfiguration
+import androidx.compose.ui.platform.LocalDensity
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.unit.Dp
+import androidx.compose.ui.unit.IntSize
+import androidx.compose.ui.unit.LayoutDirection
 import androidx.compose.ui.unit.dp
 import androidx.compose.ui.viewinterop.AndroidView
+import com.android.compose.theme.LocalAndroidColorScheme
 import com.android.systemui.communal.domain.model.CommunalContentModel
 import com.android.systemui.communal.shared.model.CommunalContentSize
 import com.android.systemui.communal.ui.viewmodel.BaseCommunalViewModel
@@ -66,21 +86,38 @@
     modifier: Modifier = Modifier,
     viewModel: BaseCommunalViewModel,
     onOpenWidgetPicker: (() -> Unit)? = null,
+    onEditDone: (() -> Unit)? = null,
 ) {
     val communalContent by viewModel.communalContent.collectAsState(initial = emptyList())
+    var removeButtonCoordinates: LayoutCoordinates? by remember { mutableStateOf(null) }
+    var toolbarSize: IntSize? by remember { mutableStateOf(null) }
+    var gridCoordinates: LayoutCoordinates? by remember { mutableStateOf(null) }
+    var isDraggingToRemove by remember { mutableStateOf(false) }
+
     Box(
         modifier = modifier.fillMaxSize().background(Color.White),
     ) {
         CommunalHubLazyGrid(
-            modifier = Modifier.height(Dimensions.GridHeight).align(Alignment.CenterStart),
+            modifier = Modifier.align(Alignment.CenterStart),
             communalContent = communalContent,
-            isEditMode = viewModel.isEditMode,
             viewModel = viewModel,
-        )
-        if (viewModel.isEditMode && onOpenWidgetPicker != null) {
-            IconButton(onClick = onOpenWidgetPicker) {
-                Icon(Icons.Default.Add, stringResource(R.string.hub_mode_add_widget_button_text))
+            contentPadding = gridContentPadding(viewModel.isEditMode, toolbarSize),
+            setGridCoordinates = { gridCoordinates = it },
+            updateDragPositionForRemove = {
+                isDraggingToRemove =
+                    checkForDraggingToRemove(it, removeButtonCoordinates, gridCoordinates)
+                isDraggingToRemove
             }
+        )
+
+        if (viewModel.isEditMode && onOpenWidgetPicker != null && onEditDone != null) {
+            Toolbar(
+                isDraggingToRemove = isDraggingToRemove,
+                setToolbarSize = { toolbarSize = it },
+                setRemoveButtonCoordinates = { removeButtonCoordinates = it },
+                onEditDone = onEditDone,
+                onOpenWidgetPicker = onOpenWidgetPicker,
+            )
         } else {
             IconButton(onClick = viewModel::onOpenWidgetEditor) {
                 Icon(Icons.Default.Edit, stringResource(R.string.button_to_open_widget_editor))
@@ -103,25 +140,38 @@
 @Composable
 private fun CommunalHubLazyGrid(
     communalContent: List<CommunalContentModel>,
-    isEditMode: Boolean,
     viewModel: BaseCommunalViewModel,
     modifier: Modifier = Modifier,
+    contentPadding: PaddingValues,
+    setGridCoordinates: (coordinates: LayoutCoordinates) -> Unit,
+    updateDragPositionForRemove: (offset: Offset) -> Boolean,
 ) {
     var gridModifier = modifier
     val gridState = rememberLazyGridState()
     var list = communalContent
     var dragDropState: GridDragDropState? = null
-    if (isEditMode && viewModel is CommunalEditModeViewModel) {
+    if (viewModel.isEditMode && viewModel is CommunalEditModeViewModel) {
         val contentListState = rememberContentListState(communalContent, viewModel)
         list = contentListState.list
-        dragDropState = rememberGridDragDropState(gridState, contentListState)
-        gridModifier = gridModifier.dragContainer(dragDropState)
+        dragDropState =
+            rememberGridDragDropState(
+                gridState = gridState,
+                contentListState = contentListState,
+                updateDragPositionForRemove = updateDragPositionForRemove
+            )
+        gridModifier =
+            gridModifier
+                .fillMaxSize()
+                .dragContainer(dragDropState, beforeContentPadding(contentPadding))
+                .onGloballyPositioned { setGridCoordinates(it) }
+    } else {
+        gridModifier = gridModifier.height(Dimensions.GridHeight)
     }
     LazyHorizontalGrid(
         modifier = gridModifier,
         state = gridState,
         rows = GridCells.Fixed(CommunalContentSize.FULL.span),
-        contentPadding = PaddingValues(horizontal = Dimensions.Spacing),
+        contentPadding = contentPadding,
         horizontalArrangement = Arrangement.spacedBy(Dimensions.Spacing),
         verticalArrangement = Arrangement.spacedBy(Dimensions.Spacing),
     ) {
@@ -130,19 +180,18 @@
             key = { index -> list[index].key },
             span = { index -> GridItemSpan(list[index].size.span) },
         ) { index ->
-            val cardModifier = Modifier.fillMaxHeight().width(Dimensions.CardWidth)
+            val cardModifier = Modifier.width(Dimensions.CardWidth)
             val size =
                 SizeF(
                     Dimensions.CardWidth.value,
                     list[index].size.dp().value,
                 )
-            if (isEditMode && dragDropState != null) {
+            if (viewModel.isEditMode && dragDropState != null) {
                 DraggableItem(dragDropState = dragDropState, enabled = true, index = index) {
                     isDragging ->
                     val elevation by animateDpAsState(if (isDragging) 4.dp else 1.dp)
                     CommunalContent(
                         modifier = cardModifier,
-                        deleteOnClick = viewModel::onDeleteWidget,
                         elevation = elevation,
                         model = list[index],
                         viewModel = viewModel,
@@ -161,6 +210,95 @@
     }
 }
 
+/**
+ * Toolbar that contains action buttons to
+ * 1) open the widget picker
+ * 2) remove a widget from the grid and
+ * 3) exit the edit mode.
+ */
+@Composable
+private fun Toolbar(
+    isDraggingToRemove: Boolean,
+    setToolbarSize: (toolbarSize: IntSize) -> Unit,
+    setRemoveButtonCoordinates: (coordinates: LayoutCoordinates) -> Unit,
+    onOpenWidgetPicker: () -> Unit,
+    onEditDone: () -> Unit,
+) {
+    Row(
+        modifier =
+            Modifier.fillMaxWidth()
+                .padding(
+                    top = Dimensions.ToolbarPaddingTop,
+                    start = Dimensions.ToolbarPaddingHorizontal,
+                    end = Dimensions.ToolbarPaddingHorizontal,
+                )
+                .onSizeChanged { setToolbarSize(it) },
+        horizontalArrangement = Arrangement.SpaceBetween,
+        verticalAlignment = Alignment.CenterVertically
+    ) {
+        val buttonContentPadding =
+            PaddingValues(
+                vertical = Dimensions.ToolbarButtonPaddingVertical,
+                horizontal = Dimensions.ToolbarButtonPaddingHorizontal,
+            )
+        val spacerModifier = Modifier.width(Dimensions.ToolbarButtonSpaceBetween)
+        Button(
+            onClick = onOpenWidgetPicker,
+            colors = filledSecondaryButtonColors(),
+            contentPadding = buttonContentPadding
+        ) {
+            Icon(Icons.Default.Add, stringResource(R.string.button_to_open_widget_editor))
+            Spacer(spacerModifier)
+            Text(
+                text = stringResource(R.string.hub_mode_add_widget_button_text),
+            )
+        }
+
+        val buttonColors =
+            if (isDraggingToRemove) filledButtonColors() else ButtonDefaults.outlinedButtonColors()
+        OutlinedButton(
+            onClick = {},
+            colors = buttonColors,
+            contentPadding = buttonContentPadding,
+            modifier = Modifier.onGloballyPositioned { setRemoveButtonCoordinates(it) },
+        ) {
+            Icon(Icons.Outlined.Delete, stringResource(R.string.button_to_open_widget_editor))
+            Spacer(spacerModifier)
+            Text(
+                text = stringResource(R.string.button_to_remove_widget),
+            )
+        }
+
+        Button(
+            onClick = onEditDone,
+            colors = filledButtonColors(),
+            contentPadding = buttonContentPadding
+        ) {
+            Text(
+                text = stringResource(R.string.hub_mode_editing_exit_button_text),
+            )
+        }
+    }
+}
+
+@Composable
+private fun filledButtonColors(): ButtonColors {
+    val colors = LocalAndroidColorScheme.current
+    return ButtonDefaults.buttonColors(
+        containerColor = colors.primary,
+        contentColor = colors.onPrimary,
+    )
+}
+
+@Composable
+private fun filledSecondaryButtonColors(): ButtonColors {
+    val colors = LocalAndroidColorScheme.current
+    return ButtonDefaults.buttonColors(
+        containerColor = colors.secondary,
+        contentColor = colors.onSecondary,
+    )
+}
+
 @Composable
 private fun CommunalContent(
     model: CommunalContentModel,
@@ -168,11 +306,9 @@
     size: SizeF,
     modifier: Modifier = Modifier,
     elevation: Dp = 0.dp,
-    deleteOnClick: ((id: Int) -> Unit)? = null,
 ) {
     when (model) {
-        is CommunalContentModel.Widget ->
-            WidgetContent(model, size, elevation, deleteOnClick, modifier)
+        is CommunalContentModel.Widget -> WidgetContent(model, size, elevation, modifier)
         is CommunalContentModel.Smartspace -> SmartspaceContent(model, modifier)
         is CommunalContentModel.Tutorial -> TutorialContent(modifier)
         is CommunalContentModel.Umo -> Umo(viewModel, modifier)
@@ -184,19 +320,12 @@
     model: CommunalContentModel.Widget,
     size: SizeF,
     elevation: Dp,
-    deleteOnClick: ((id: Int) -> Unit)?,
     modifier: Modifier = Modifier,
 ) {
-    // TODO(b/309009246): update background color
     Card(
-        modifier = modifier.fillMaxSize().background(Color.White),
+        modifier = modifier.height(size.height.dp),
         elevation = CardDefaults.cardElevation(draggedElevation = elevation),
     ) {
-        if (deleteOnClick != null) {
-            IconButton(onClick = { deleteOnClick(model.appWidgetId) }) {
-                Icon(Icons.Default.Close, stringResource(R.string.button_to_remove_widget))
-            }
-        }
         AndroidView(
             modifier = modifier,
             factory = { context ->
@@ -249,6 +378,60 @@
     )
 }
 
+/**
+ * Returns the `contentPadding` of the grid. Use the vertical padding to push the grid content area
+ * below the toolbar and let the grid take the max size. This ensures the item can be dragged
+ * outside the grid over the toolbar, without part of it getting clipped by the container.
+ */
+@Composable
+private fun gridContentPadding(isEditMode: Boolean, toolbarSize: IntSize?): PaddingValues {
+    if (!isEditMode || toolbarSize == null) {
+        return PaddingValues(horizontal = Dimensions.Spacing)
+    }
+    val configuration = LocalConfiguration.current
+    val density = LocalDensity.current
+    val screenHeight = configuration.screenHeightDp.dp
+    val toolbarHeight = with(density) { Dimensions.ToolbarPaddingTop + toolbarSize.height.toDp() }
+    val verticalPadding =
+        ((screenHeight - toolbarHeight - Dimensions.GridHeight) / 2).coerceAtLeast(
+            Dimensions.Spacing
+        )
+    return PaddingValues(
+        start = Dimensions.ToolbarPaddingHorizontal,
+        end = Dimensions.ToolbarPaddingHorizontal,
+        top = verticalPadding + toolbarHeight,
+        bottom = verticalPadding
+    )
+}
+
+@Composable
+private fun beforeContentPadding(paddingValues: PaddingValues): ContentPaddingInPx {
+    return with(LocalDensity.current) {
+        ContentPaddingInPx(
+            startPadding = paddingValues.calculateLeftPadding(LayoutDirection.Ltr).toPx(),
+            topPadding = paddingValues.calculateTopPadding().toPx()
+        )
+    }
+}
+
+/**
+ * Check whether the pointer position that the item is being dragged at is within the coordinates of
+ * the remove button in the toolbar. Returns true if the item is removable.
+ */
+private fun checkForDraggingToRemove(
+    offset: Offset,
+    removeButtonCoordinates: LayoutCoordinates?,
+    gridCoordinates: LayoutCoordinates?,
+): Boolean {
+    if (removeButtonCoordinates == null || gridCoordinates == null) {
+        return false
+    }
+    val pointer = gridCoordinates.positionInWindow() + offset
+    val removeButton = removeButtonCoordinates.positionInWindow()
+    return pointer.x in removeButton.x..removeButton.x + removeButtonCoordinates.size.width &&
+        pointer.y in removeButton.y..removeButton.y + removeButtonCoordinates.size.height
+}
+
 private fun CommunalContentSize.dp(): Dp {
     return when (this) {
         CommunalContentSize.FULL -> Dimensions.CardHeightFull
@@ -257,6 +440,8 @@
     }
 }
 
+data class ContentPaddingInPx(val startPadding: Float, val topPadding: Float)
+
 object Dimensions {
     val CardWidth = 464.dp
     val CardHeightFull = 630.dp
@@ -264,4 +449,11 @@
     val CardHeightThird = 199.dp
     val GridHeight = CardHeightFull
     val Spacing = 16.dp
+
+    // The sizing/padding of the toolbar in glanceable hub edit mode
+    val ToolbarPaddingTop = 27.dp
+    val ToolbarPaddingHorizontal = 16.dp
+    val ToolbarButtonPaddingHorizontal = 24.dp
+    val ToolbarButtonPaddingVertical = 16.dp
+    val ToolbarButtonSpaceBetween = 8.dp
 }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/GridDragDropState.kt b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/GridDragDropState.kt
index 6cfa2f2..5451d05 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/GridDragDropState.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/GridDragDropState.kt
@@ -48,12 +48,18 @@
 @Composable
 fun rememberGridDragDropState(
     gridState: LazyGridState,
-    contentListState: ContentListState
+    contentListState: ContentListState,
+    updateDragPositionForRemove: (offset: Offset) -> Boolean,
 ): GridDragDropState {
     val scope = rememberCoroutineScope()
     val state =
         remember(gridState, contentListState) {
-            GridDragDropState(state = gridState, contentListState = contentListState, scope = scope)
+            GridDragDropState(
+                state = gridState,
+                contentListState = contentListState,
+                scope = scope,
+                updateDragPositionForRemove = updateDragPositionForRemove
+            )
         }
     LaunchedEffect(state) {
         while (true) {
@@ -67,23 +73,30 @@
 /**
  * Handles drag and drop cards in the glanceable hub. While dragging to move, other items that are
  * affected will dynamically get positioned and the state is tracked by [ContentListState]. When
- * dragging to remove, affected cards will be moved and [ContentListState.onRemove] is called to
- * remove the dragged item. On dragging ends, call [ContentListState.onSaveList] to persist the
- * change.
+ * dragging to remove, affected cards will be moved and [updateDragPositionForRemove] is called to
+ * check whether the dragged item can be removed. On dragging ends, call [ContentListState.onRemove]
+ * to remove the dragged item if condition met and call [ContentListState.onSaveList] to persist any
+ * change in ordering.
  */
 class GridDragDropState
 internal constructor(
     private val state: LazyGridState,
     private val contentListState: ContentListState,
     private val scope: CoroutineScope,
+    private val updateDragPositionForRemove: (offset: Offset) -> Boolean
 ) {
     var draggingItemIndex by mutableStateOf<Int?>(null)
         private set
 
+    var isDraggingToRemove by mutableStateOf(false)
+        private set
+
     internal val scrollChannel = Channel<Float>()
 
     private var draggingItemDraggedDelta by mutableStateOf(Offset.Zero)
     private var draggingItemInitialOffset by mutableStateOf(Offset.Zero)
+    private var dragStartPointerOffset by mutableStateOf(Offset.Zero)
+
     internal val draggingItemOffset: Offset
         get() =
             draggingItemLayoutInfo?.let { item ->
@@ -94,27 +107,36 @@
     private val draggingItemLayoutInfo: LazyGridItemInfo?
         get() = state.layoutInfo.visibleItemsInfo.firstOrNull { it.index == draggingItemIndex }
 
-    internal fun onDragStart(offset: Offset) {
+    internal fun onDragStart(offset: Offset, contentOffset: Offset) {
         state.layoutInfo.visibleItemsInfo
             .firstOrNull { item ->
+                // grid item offset is based off grid content container so we need to deduct
+                // before content padding from the initial pointer position
                 item.isEditable &&
-                    offset.x.toInt() in item.offset.x..item.offsetEnd.x &&
-                    offset.y.toInt() in item.offset.y..item.offsetEnd.y
+                    (offset.x - contentOffset.x).toInt() in item.offset.x..item.offsetEnd.x &&
+                    (offset.y - contentOffset.y).toInt() in item.offset.y..item.offsetEnd.y
             }
             ?.apply {
+                dragStartPointerOffset = offset - this.offset.toOffset()
                 draggingItemIndex = index
                 draggingItemInitialOffset = this.offset.toOffset()
             }
     }
 
     internal fun onDragInterrupted() {
-        if (draggingItemIndex != null) {
+        draggingItemIndex?.let {
+            if (isDraggingToRemove) {
+                contentListState.onRemove(it)
+                isDraggingToRemove = false
+                updateDragPositionForRemove(Offset.Zero)
+            }
             // persist list editing changes on dragging ends
             contentListState.onSaveList()
             draggingItemIndex = null
         }
         draggingItemDraggedDelta = Offset.Zero
         draggingItemInitialOffset = Offset.Zero
+        dragStartPointerOffset = Offset.Zero
     }
 
     internal fun onDrag(offset: Offset) {
@@ -152,18 +174,13 @@
                 contentListState.onMove(draggingItem.index, targetItem.index)
             }
             draggingItemIndex = targetItem.index
+            isDraggingToRemove = false
         } else {
             val overscroll = checkForOverscroll(startOffset, endOffset)
             if (overscroll != 0f) {
                 scrollChannel.trySend(overscroll)
             }
-            val removeOffset = checkForRemove(startOffset)
-            if (removeOffset != 0f) {
-                draggingItemIndex?.let {
-                    contentListState.onRemove(it)
-                    draggingItemIndex = null
-                }
-            }
+            isDraggingToRemove = checkForRemove(startOffset)
         }
     }
 
@@ -185,14 +202,11 @@
         }
     }
 
-    // TODO(b/309968801): a temporary solution to decide whether to remove card when it's dragged up
-    //  and out of grid. Once we have a taskbar, calculate the intersection of the dragged item with
-    //  the Remove button.
-    private fun checkForRemove(startOffset: Offset): Float {
+    /** Calls the callback with the updated drag position and returns whether to remove the item. */
+    private fun checkForRemove(startOffset: Offset): Boolean {
         return if (draggingItemDraggedDelta.y < 0)
-            (startOffset.y + Dimensions.CardHeightHalf.value - state.layoutInfo.viewportStartOffset)
-                .coerceAtMost(0f)
-        else 0f
+            updateDragPositionForRemove(startOffset + dragStartPointerOffset)
+        else false
     }
 }
 
@@ -204,14 +218,22 @@
     return Offset(x + size.width, y + size.height)
 }
 
-fun Modifier.dragContainer(dragDropState: GridDragDropState): Modifier {
-    return pointerInput(dragDropState) {
+fun Modifier.dragContainer(
+    dragDropState: GridDragDropState,
+    beforeContentPadding: ContentPaddingInPx
+): Modifier {
+    return pointerInput(dragDropState, beforeContentPadding) {
         detectDragGesturesAfterLongPress(
             onDrag = { change, offset ->
                 change.consume()
                 dragDropState.onDrag(offset = offset)
             },
-            onDragStart = { offset -> dragDropState.onDragStart(offset) },
+            onDragStart = { offset ->
+                dragDropState.onDragStart(
+                    offset,
+                    Offset(beforeContentPadding.startPadding, beforeContentPadding.topPadding)
+                )
+            },
             onDragEnd = { dragDropState.onDragInterrupted() },
             onDragCancel = { dragDropState.onDragInterrupted() }
         )
@@ -237,6 +259,7 @@
             Modifier.zIndex(1f).graphicsLayer {
                 translationX = dragDropState.draggingItemOffset.x
                 translationY = dragDropState.draggingItemOffset.y
+                alpha = if (dragDropState.isDraggingToRemove) 0.5f else 1f
             }
         } else {
             Modifier.animateItemPlacement()
diff --git a/packages/SystemUI/res/values/strings.xml b/packages/SystemUI/res/values/strings.xml
index 78b701c..13d2fea 100644
--- a/packages/SystemUI/res/values/strings.xml
+++ b/packages/SystemUI/res/values/strings.xml
@@ -1066,9 +1066,11 @@
     <!-- Description for the button that opens the widget editor on click. [CHAR LIMIT=50] -->
     <string name="button_to_open_widget_editor">Open the widget editor</string>
     <!-- Description for the button that removes a widget on click. [CHAR LIMIT=50] -->
-    <string name="button_to_remove_widget">Remove a widget</string>
+    <string name="button_to_remove_widget">Remove</string>
     <!-- Text for the button that launches the hub mode widget picker. [CHAR LIMIT=50] -->
-    <string name="hub_mode_add_widget_button_text">Add Widget</string>
+    <string name="hub_mode_add_widget_button_text">Add widget</string>
+    <!-- Text for the button that exits the hub mode editing mode. [CHAR LIMIT=50] -->
+    <string name="hub_mode_editing_exit_button_text">Done</string>
 
     <!-- Related to user switcher --><skip/>
 
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 7b94fc1..573a748 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt
@@ -76,6 +76,10 @@
                     Intent(applicationContext, WidgetPickerActivity::class.java)
                 )
             },
+            onEditDone = {
+                // TODO(b/315154364): in a separate change, lock the device and transition to GH
+                finish()
+            }
         )
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/compose/BaseComposeFacade.kt b/packages/SystemUI/src/com/android/systemui/compose/BaseComposeFacade.kt
index 65d4495..3a92739 100644
--- a/packages/SystemUI/src/com/android/systemui/compose/BaseComposeFacade.kt
+++ b/packages/SystemUI/src/com/android/systemui/compose/BaseComposeFacade.kt
@@ -63,6 +63,7 @@
         activity: ComponentActivity,
         viewModel: BaseCommunalViewModel,
         onOpenWidgetPicker: () -> Unit,
+        onEditDone: () -> Unit,
     )
 
     /** Create a [View] to represent [viewModel] on screen. */