package org.ambraproject.rhino.view;

import com.google.common.collect.ImmutableList;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;
import com.google.gson.JsonSerializationContext;
import org.ambraproject.models.Annotation;
import org.ambraproject.models.Article;
import org.ambraproject.rhino.view.article.ArticleVisibility;

import java.util.Collection;
import java.util.Comparator;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

/**
 * View of a comment, with nested views of all its children.
 */
public class AnnotationOutputView implements JsonOutputView {

  private final ArticleVisibility parentArticle;
  private final Annotation comment;
  private final Date mostRecentActivity;
  private final ImmutableList<AnnotationOutputView> replies;

  private AnnotationOutputView(ArticleVisibility parentArticle,
                               Annotation comment,
                               Date mostRecentActivity,
                               List<AnnotationOutputView> replies) {
    this.parentArticle = Objects.requireNonNull(parentArticle);
    this.comment = Objects.requireNonNull(comment);
    this.mostRecentActivity = Objects.requireNonNull(mostRecentActivity);
    this.replies = ImmutableList.copyOf(replies);
  }

  public static class Factory {
    private final ArticleVisibility parentArticle;
    private final Map<Long, List<Annotation>> commentsByParent;

    /**
     * @param parentArticle an article
     * @param comments      all comments belonging to the parent article
     */
    public Factory(Article parentArticle, Collection<Annotation> comments) {
      this.parentArticle = ArticleVisibility.create(parentArticle);
      this.commentsByParent = comments.stream()
          .filter(comment -> comment.getParentID() != null)
          .collect(Collectors.groupingBy(Annotation::getParentID));
    }

    /**
     * @param comment a comment belonging to this object's parent article
     * @return a view of the comment and all its children
     */
    public AnnotationOutputView buildView(Annotation comment) {
      List<Annotation> childObjects = commentsByParent.getOrDefault(comment.getID(), ImmutableList.of());
      List<AnnotationOutputView> childViews = childObjects.stream()
          .sorted(BY_DATE)
          .map(this::buildView) // recursion (terminal case is when childObjects is empty)
          .collect(Collectors.toList());
      Date mostRecentActivity = findMostRecentActivity(comment, childViews);
      return new AnnotationOutputView(parentArticle, comment, mostRecentActivity, childViews);
    }

    private static Date findMostRecentActivity(Annotation comment, Collection<AnnotationOutputView> childViews) {
      return childViews.stream()
          .map(view -> view.mostRecentActivity)
          .max(Comparator.naturalOrder())
          .orElse(comment.getCreated());
    }
  }

  public static final Comparator<Annotation> BY_DATE = Comparator.comparing(Annotation::getCreated);

  @Override
  public JsonElement serialize(JsonSerializationContext context) {
    JsonObject serialized = context.serialize(comment).getAsJsonObject();
    serialized.remove("articleID");
    serialized.add("parentArticle", context.serialize(parentArticle));
    serialized.add("mostRecentActivity", context.serialize(mostRecentActivity));
    serialized.add("replies", context.serialize(replies));
    return serialized;
  }

}

@@ -31,16 +31,61 @@
 .view;%0A%0A
+import com.google.common.base.Preconditions;%0A
 import c
@@ -800,49 +800,8 @@
 nt;%0A
-  private final Date mostRecentActivity;%0A
   pr
@@ -858,16 +858,93 @@
 plies;%0A%0A
+  private final int replyTreeSize;%0A  private final Date mostRecentActivity;%0A%0A
   privat
@@ -1085,98 +1085,117 @@
     
-Date mostRecentActivity,%0A                               List%3CAnnotationOutputView%3E replies
+List%3CAnnotationOutputView%3E replies,%0A                               int replyTreeSize, Date mostRecentActivity
 ) %7B%0A
@@ -1310,16 +1310,160 @@
 mment);%0A
+    this.replies = ImmutableList.copyOf(replies);%0A%0A    Preconditions.checkArgument(replyTreeSize %3E= 0);%0A    this.replyTreeSize = replyTreeSize;%0A
     this
@@ -1532,58 +1532,8 @@
 y);%0A
-    this.replies = ImmutableList.copyOf(replies);%0A
   %7D%0A
@@ -2677,24 +2677,87 @@
 s.toList());
+%0A%0A      int replyTreeSize = calculateReplyTreeSize(childViews);
 %0A      Date 
@@ -2818,24 +2818,25 @@
 hildViews);%0A
+%0A
       return
@@ -2889,38 +2889,250 @@
 nt, 
-mostRecentActivity, childViews
+childViews, replyTreeSize, mostRecentActivity);%0A    %7D%0A%0A    private static int calculateReplyTreeSize(Collection%3CAnnotationOutputView%3E childViews) %7B%0A      return childViews.size() + childViews.stream().mapToInt(view -%3E view.replyTreeSize).sum(
 );%0A 
@@ -3784,24 +3784,95 @@
 tArticle));%0A
+    serialized.add(%22replyTreeSize%22, context.serialize(replyTreeSize));%0A
     serializ

Add replyTreeSize attribute to comments