Fix and cleanup smart replies checking

1. We now apply smart replies to the first action button that has
   (freeform) remote input. For example, if the notification have
   both reply and reply all buttons, smart replies will be applied
   to the first button only.
2. Enforced getAllowGeneratedReplies check in system generated
   smart replies.
3. Fixed an bug that smartRepliesAdded is not called for system
   generated smart replies.

BUG: 111546109
BUG: 111406942

Test: atest com.android.server.notification.NotificationTest
Test: Check apps generated smart replies via Messenger
Test: Check system generated smart replies via Hangouts

Change-Id: I4db34f557f7e9988be612e4162347b86393d1faf

3 files changed, 159 insertions, 43 deletions

diff --git a/core/java/android/app/Notification.java b/core/java/android/app/Notification.java
index 3638bc48d2b5..01328dc29353 100644
--- a/core/java/android/app/Notification.java
+++ b/core/java/android/app/Notification.java
@@ -68,6 +68,7 @@ import android.text.style.RelativeSizeSpan;
 import android.text.style.TextAppearanceSpan;
 import android.util.ArraySet;
 import android.util.Log;
+import android.util.Pair;
 import android.util.SparseArray;
 import android.util.proto.ProtoOutputStream;
 import android.view.Gravity;
@@ -3129,6 +3130,37 @@ public class Notification implements Parcelable
         return false;
     }
 
+
+    /**
+     * Finds and returns a remote input and its corresponding action.
+     *
+     * @param requiresFreeform requires the remoteinput to allow freeform or not.
+     * @return the result pair, {@code null} if no result is found.
+     *
+     * @hide
+     */
+    @Nullable
+    public Pair<RemoteInput, Action> findRemoteInputActionPair(boolean requiresFreeform) {
+        if (actions == null) {
+            return null;
+        }
+        for (Notification.Action action : actions) {
+            if (action.getRemoteInputs() == null) {
+                continue;
+            }
+            RemoteInput resultRemoteInput = null;
+            for (RemoteInput remoteInput : action.getRemoteInputs()) {
+                if (remoteInput.getAllowFreeFormInput() || !requiresFreeform) {
+                    resultRemoteInput = remoteInput;
+                }
+            }
+            if (resultRemoteInput != null) {
+                return Pair.create(resultRemoteInput, action);
+            }
+        }
+        return null;
+    }
+
     /**
      * Builder class for {@link Notification} objects.
      *
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentView.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentView.java
index 4963a0c6ae01..4ef8dbb19318 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentView.java
@@ -26,6 +26,7 @@ import android.service.notification.StatusBarNotification;
 import android.util.ArraySet;
 import android.util.AttributeSet;
 import android.util.Log;
+import android.util.Pair;
 import android.view.MotionEvent;
 import android.view.NotificationHeaderView;
 import android.view.View;
@@ -40,12 +41,12 @@ import com.android.internal.util.ArrayUtils;
 import com.android.internal.util.ContrastColorUtil;
 import com.android.systemui.Dependency;
 import com.android.systemui.R;
-import com.android.systemui.statusbar.notification.NotificationData;
 import com.android.systemui.statusbar.RemoteInputController;
 import com.android.systemui.statusbar.SmartReplyController;
 import com.android.systemui.statusbar.TransformableView;
-import com.android.systemui.statusbar.notification.row.wrapper.NotificationCustomViewWrapper;
+import com.android.systemui.statusbar.notification.NotificationData;
 import com.android.systemui.statusbar.notification.NotificationUtils;
+import com.android.systemui.statusbar.notification.row.wrapper.NotificationCustomViewWrapper;
 import com.android.systemui.statusbar.notification.row.wrapper.NotificationViewWrapper;
 import com.android.systemui.statusbar.phone.NotificationGroupManager;
 import com.android.systemui.statusbar.policy.RemoteInputView;
@@ -1188,6 +1189,7 @@ public class NotificationContentView extends FrameLayout {
         updateSingleLineView();
         updateAmbientSingleLineView();
     }
+
     private void updateSingleLineView() {
         if (mIsChildInGroup) {
             boolean isNewView = mSingleLineView == null;
@@ -1223,53 +1225,44 @@ public class NotificationContentView extends FrameLayout {
             return;
         }
 
-        boolean enableSmartReplies = (mSmartReplyConstants.isEnabled()
+        Notification notification = entry.notification.getNotification();
+
+        Pair<RemoteInput, Notification.Action> remoteInputActionPair =
+                entry.notification.getNotification().findRemoteInputActionPair(false /*freeform */);
+        Pair<RemoteInput, Notification.Action> freeformRemoteInputActionPair =
+                notification.findRemoteInputActionPair(true /*freeform */);
+
+        boolean enableAppGeneratedSmartReplies = (mSmartReplyConstants.isEnabled()
                 && (!mSmartReplyConstants.requiresTargetingP()
-                    || entry.targetSdk >= Build.VERSION_CODES.P));
+                || entry.targetSdk >= Build.VERSION_CODES.P));
 
-        boolean hasRemoteInput = false;
         RemoteInput remoteInputWithChoices = null;
-        PendingIntent pendingIntentWithChoices = null;
+        PendingIntent pendingIntentWithChoices= null;
         CharSequence[] choices = null;
-
-        Notification.Action[] actions = entry.notification.getNotification().actions;
-        if (actions != null) {
-            for (Notification.Action a : actions) {
-                if (a.getRemoteInputs() == null) {
-                    continue;
-                }
-                for (RemoteInput ri : a.getRemoteInputs()) {
-                    boolean showRemoteInputView = ri.getAllowFreeFormInput();
-                    boolean showSmartReplyView = enableSmartReplies
-                            && (ArrayUtils.isEmpty(ri.getChoices())
-                            || (showRemoteInputView && !ArrayUtils.isEmpty(entry.smartReplies)));
-                    if (showRemoteInputView) {
-                        hasRemoteInput = true;
-                    }
-                    if (showSmartReplyView) {
-                        remoteInputWithChoices = ri;
-                        pendingIntentWithChoices = a.actionIntent;
-                        if (!ArrayUtils.isEmpty(ri.getChoices())) {
-                            choices = ri.getChoices();
-                        } else {
-                            choices = entry.smartReplies;
-                        }
-                    }
-                    if (showRemoteInputView || showSmartReplyView) {
-                        break;
-                    }
-                }
-            }
-        }
-
-        applyRemoteInput(entry, hasRemoteInput);
+        if (enableAppGeneratedSmartReplies
+                && remoteInputActionPair != null
+                && !ArrayUtils.isEmpty(remoteInputActionPair.first.getChoices())) {
+            // app generated smart replies
+            remoteInputWithChoices = remoteInputActionPair.first;
+            pendingIntentWithChoices = remoteInputActionPair.second.actionIntent;
+            choices = remoteInputActionPair.first.getChoices();
+        } else if (!ArrayUtils.isEmpty(entry.smartReplies)
+                && freeformRemoteInputActionPair != null
+                && freeformRemoteInputActionPair.second.getAllowGeneratedReplies()) {
+            // system generated smart replies
+            remoteInputWithChoices = freeformRemoteInputActionPair.first;
+            pendingIntentWithChoices = freeformRemoteInputActionPair.second.actionIntent;
+            choices = entry.smartReplies;
+        }
+
+        applyRemoteInput(entry, freeformRemoteInputActionPair != null);
         applySmartReplyView(remoteInputWithChoices, pendingIntentWithChoices, entry, choices);
     }
 
-    private void applyRemoteInput(NotificationData.Entry entry, boolean hasRemoteInput) {
+    private void applyRemoteInput(NotificationData.Entry entry, boolean hasFreeformRemoteInput) {
         View bigContentView = mExpandedChild;
         if (bigContentView != null) {
-            mExpandedRemoteInput = applyRemoteInput(bigContentView, entry, hasRemoteInput,
+            mExpandedRemoteInput = applyRemoteInput(bigContentView, entry, hasFreeformRemoteInput,
                     mPreviousExpandedRemoteInputIntent, mCachedExpandedRemoteInput,
                     mExpandedWrapper);
         } else {
@@ -1284,7 +1277,8 @@ public class NotificationContentView extends FrameLayout {
 
         View headsUpContentView = mHeadsUpChild;
         if (headsUpContentView != null) {
-            mHeadsUpRemoteInput = applyRemoteInput(headsUpContentView, entry, hasRemoteInput,
+            mHeadsUpRemoteInput = applyRemoteInput(
+                    headsUpContentView, entry, hasFreeformRemoteInput,
                     mPreviousHeadsUpRemoteInputIntent, mCachedHeadsUpRemoteInput, mHeadsUpWrapper);
         } else {
             mHeadsUpRemoteInput = null;
@@ -1370,8 +1364,8 @@ public class NotificationContentView extends FrameLayout {
             mExpandedSmartReplyView =
                     applySmartReplyView(mExpandedChild, remoteInput, pendingIntent, entry, choices);
             if (mExpandedSmartReplyView != null && remoteInput != null
-                    && remoteInput.getChoices() != null && remoteInput.getChoices().length > 0) {
-                mSmartReplyController.smartRepliesAdded(entry, remoteInput.getChoices().length);
+                    && choices != null && choices.length > 0) {
+                mSmartReplyController.smartRepliesAdded(entry, choices.length);
             }
         }
     }
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/NotificationTest.java b/services/tests/uiservicestests/src/com/android/server/notification/NotificationTest.java
index 9db823c0d3f5..55c16b367fab 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationTest.java
@@ -48,6 +48,7 @@ import android.support.test.runner.AndroidJUnit4;
 import android.text.SpannableStringBuilder;
 import android.text.Spanned;
 import android.text.style.StyleSpan;
+import android.util.Pair;
 import android.widget.RemoteViews;
 
 import com.android.server.UiServiceTestCase;
@@ -541,5 +542,94 @@ public class NotificationTest extends UiServiceTestCase {
 
         assertFalse(Notification.areActionsVisiblyDifferent(n1, n2));
     }
+
+    @Test
+    public void testFreeformRemoteInputActionPair_noRemoteInput() {
+        PendingIntent intent = mock(PendingIntent.class);
+        Icon icon = mock(Icon.class);
+        Notification notification = new Notification.Builder(mContext, "test")
+                .addAction(new Notification.Action.Builder(icon, "TEXT 1", intent)
+                        .build())
+                .build();
+        assertNull(notification.findRemoteInputActionPair(false));
+    }
+
+    @Test
+    public void testFreeformRemoteInputActionPair_hasRemoteInput() {
+        PendingIntent intent = mock(PendingIntent.class);
+        Icon icon = mock(Icon.class);
+
+        RemoteInput remoteInput = new RemoteInput.Builder("a").build();
+
+        Notification.Action actionWithRemoteInput =
+                new Notification.Action.Builder(icon, "TEXT 1", intent)
+                        .addRemoteInput(remoteInput)
+                        .addRemoteInput(remoteInput)
+                        .build();
+
+        Notification.Action actionWithoutRemoteInput =
+                new Notification.Action.Builder(icon, "TEXT 2", intent)
+                        .build();
+
+        Notification notification = new Notification.Builder(mContext, "test")
+                .addAction(actionWithoutRemoteInput)
+                .addAction(actionWithRemoteInput)
+                .build();
+
+        Pair<RemoteInput, Notification.Action> remoteInputActionPair =
+                notification.findRemoteInputActionPair(false);
+
+        assertNotNull(remoteInputActionPair);
+        assertEquals(remoteInput, remoteInputActionPair.first);
+        assertEquals(actionWithRemoteInput, remoteInputActionPair.second);
+    }
+
+    @Test
+    public void testFreeformRemoteInputActionPair_requestFreeform_noFreeformRemoteInput() {
+        PendingIntent intent = mock(PendingIntent.class);
+        Icon icon = mock(Icon.class);
+        Notification notification = new Notification.Builder(mContext, "test")
+                .addAction(new Notification.Action.Builder(icon, "TEXT 1", intent)
+                        .addRemoteInput(
+                                new RemoteInput.Builder("a")
+                                        .setAllowFreeFormInput(false).build())
+                        .build())
+                .build();
+        assertNull(notification.findRemoteInputActionPair(true));
+    }
+
+    @Test
+    public void testFreeformRemoteInputActionPair_requestFreeform_hasFreeformRemoteInput() {
+        PendingIntent intent = mock(PendingIntent.class);
+        Icon icon = mock(Icon.class);
+
+        RemoteInput remoteInput =
+                new RemoteInput.Builder("a").setAllowFreeFormInput(false).build();
+        RemoteInput freeformRemoteInput =
+                new RemoteInput.Builder("b").setAllowFreeFormInput(true).build();
+
+        Notification.Action actionWithFreeformRemoteInput =
+                new Notification.Action.Builder(icon, "TEXT 1", intent)
+                        .addRemoteInput(remoteInput)
+                        .addRemoteInput(freeformRemoteInput)
+                        .build();
+
+        Notification.Action actionWithoutFreeformRemoteInput =
+                new Notification.Action.Builder(icon, "TEXT 2", intent)
+                        .addRemoteInput(remoteInput)
+                        .build();
+
+        Notification notification = new Notification.Builder(mContext, "test")
+                .addAction(actionWithoutFreeformRemoteInput)
+                .addAction(actionWithFreeformRemoteInput)
+                .build();
+
+        Pair<RemoteInput, Notification.Action> remoteInputActionPair =
+                notification.findRemoteInputActionPair(true);
+
+        assertNotNull(remoteInputActionPair);
+        assertEquals(freeformRemoteInput, remoteInputActionPair.first);
+        assertEquals(actionWithFreeformRemoteInput, remoteInputActionPair.second);
+    }
 }