path: root/src/com/google/devrel/cluestick/studioclient/SearchResultsTree.java

/*
 * Copyright (C) 2016 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.google.devrel.cluestick.studioclient;

import com.appspot.cluestick_server.search.model.CodeResult;
import com.appspot.cluestick_server.search.model.Result;

import com.intellij.icons.AllIcons;
import com.intellij.ui.ColoredTreeCellRenderer;
import com.intellij.ui.SimpleTextAttributes;
import com.intellij.ui.treeStructure.Tree;
import com.intellij.util.ui.UIUtil;

import org.jetbrains.annotations.NotNull;

import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import javax.swing.Icon;
import javax.swing.JTree;
import javax.swing.tree.DefaultMutableTreeNode;
import javax.swing.tree.TreeNode;
import javax.swing.tree.TreePath;

/**
 * Tree structure containing search results for display in a JTree. Results are displayed as a
 * collection of repos. This tree is not expanded by default.
 */
public final class SearchResultsTree extends Tree {
  private static final SimpleTextAttributes SUFFIX_ATTRIBUTES = new SimpleTextAttributes(
      SimpleTextAttributes.STYLE_ITALIC, UIUtil.getInactiveTextColor());

  private final LayoutNode symbolHeading;
  private final LayoutNode resultHeading;

  public SearchResultsTree(Symbol symbol, List<Result> results) {
    super(new DefaultMutableTreeNode());
    DefaultMutableTreeNode top = (DefaultMutableTreeNode) getModel().getRoot();

    setRootVisible(false);
    setCellRenderer(new CellRenderer());

    symbolHeading = new LayoutNode();
    symbolHeading.text = "Symbol";
    symbolHeading.isHeading = true;
    top.add(symbolHeading);

    LayoutNode symbolNode = new LayoutNode();
    if (symbol.packageName != null && !symbol.packageName.isEmpty()) {
      symbolNode.prefix = symbol.packageName + ".";
    }
    symbolNode.icon = AllIcons.Nodes.EjbFinderMethod;
    symbolNode.text = symbol.symbolName;
    symbolHeading.add(symbolNode);

    resultHeading = new LayoutNode();
    resultHeading.text = "Found results";
    resultHeading.count = results.size();
    resultHeading.isHeading = true;
    top.add(resultHeading);

    addResults(resultHeading, results);
  }

  /**
   * Determines whether the {@link Result} has code.
   * @param result The result to check.
   * @return Whether the result safely has code.
   */
  public static boolean hasCode(Result result) {
    return result.getCode() != null && result.getCode().getLines() != null;
  }

  private void addResults(DefaultMutableTreeNode parent, List<Result> results) {
    Map<String, RepoNode> repoNodes = new LinkedHashMap<String, RepoNode>();  // linked provides insert ordering

    for (Result result : results) {
      if (!hasCode(result)) {
        parent.add(new ResultNode(result));
        continue;
      }

      // If this is a code node, then group by repo/branch pair. Retrieve or create the RepoNode
      // for this pair, add and update its total count.
      CodeResult code = result.getCode();
      String key = String.format("%s:%s", code.getRepo(), code.getBranch());
      RepoNode repoNode = repoNodes.get(key);
      if (repoNode == null) {
        repoNode = new RepoNode(code.getRepo(), code.getBranch());
        repoNodes.put(key, repoNode);
      }
      repoNode.add(new ResultNode(result));
    }

    // Now, add the repo nodes in order of first seen.
    for (RepoNode repoNode : repoNodes.values()) {
      parent.add(repoNode);
    }
  }

  /**
   * RepoNode represents a whole repository (aka, repo/branch).
   */
  private static class RepoNode extends DefaultMutableTreeNode implements SearchResultsNode {
    private final String repo;
    private final String branch;
    private final ResultUtils.GitHubInfo info;

    RepoNode(String repo, String branch) {
      this.repo = repo;
      this.branch = branch;

      info = ResultUtils.parseRepoName(repo);
    }

    public Icon getIcon() {
      if (info != null) {
        return IconFetcher.GitHub;
      }
      return null;
    }

    public String getSuffix() {
      if (this.branch == null || this.branch.isEmpty() || this.branch.equals("master")) {
        return null;
      }
      return this.branch;
    }

    public String getPrefix() {
      if (info != null) {
        return info.user + "/";
      }
      return null;
    }

    @Override
    public String getURL() {
      if (info != null) {
        // TODO(thorogood): Deal with non-master branch.
        return String.format("https://github.com/%s/%s", info.user, info.repo);
      }
      return null;
    }

    @Override
    public Result getSearchResult() {
      return null;
    }

    @Override
    public String toString() {
      if (info != null) {
        return info.repo;
      }
      return this.repo;
    }
  }


  /**
   * LayoutNode is a generic node for use inside SearchResultsTree.
   */
  private static class LayoutNode extends DefaultMutableTreeNode {
    String prefix;
    String suffix;
    String text;
    Icon icon;
    boolean isHeading;
    int count = -1;

    @NotNull
    public String getText() {
      if (text == null) {
        return "?";
      }
      return text;
    }
  }

  /**
   * Node representing a {@link Result}.
   */
  public static class ResultNode extends DefaultMutableTreeNode implements SearchResultsNode {
    public ResultNode(Result result) {
      super(result);
    }

    @Override
    public String getURL() {
      return getSearchResult().getUrl();
    }

    @Override
    public Result getSearchResult() {
      return (Result) getUserObject();
    }
  }

  private static class CellRenderer extends ColoredTreeCellRenderer {
    @Override
    public void customizeCellRenderer(JTree tree, Object value,
        boolean selected, boolean expanded, boolean leaf,
        int row, boolean hasFocus) {

      if (value instanceof LayoutNode) {
        LayoutNode node = (LayoutNode) value;
        setIcon(node.icon);

        if (node.prefix != null) {
          append(node.prefix, SimpleTextAttributes.GRAYED_ATTRIBUTES, false);
        }

        SimpleTextAttributes defaultAttributes = SimpleTextAttributes.REGULAR_ATTRIBUTES;
        if (node.isHeading) {
          defaultAttributes = SimpleTextAttributes.REGULAR_BOLD_ATTRIBUTES;
        }
        append(node.getText(), defaultAttributes, true);

        if (node.suffix != null) {
          append(" " + node.suffix, SUFFIX_ATTRIBUTES, false);
        }
        if (node.count >= 0) {
          String format = " (%d result" + (node.count != 1 ? "s" : "") + ")";
          append(String.format(format, node.count), SUFFIX_ATTRIBUTES, false);
        }
        return;
      }

      if (value instanceof RepoNode) {
        RepoNode node = (RepoNode) value;
        setIcon(node.getIcon());

        String prefix = node.getPrefix();
        if (prefix != null) {
          append(prefix, SimpleTextAttributes.GRAYED_ATTRIBUTES, false);
        }

        append(node.toString(), SimpleTextAttributes.REGULAR_ATTRIBUTES, true);

        String suffix = node.getSuffix();
        if (suffix != null) {
          append(" " + suffix, SimpleTextAttributes.GRAYED_ATTRIBUTES, false);
        }

        return;
      }

      if (value instanceof ResultNode) {
        ResultNode node = (ResultNode) value;
        Result result = node.getSearchResult();

        if (!hasCode(result)) {
          setIcon(AllIcons.Actions.CreateFromUsage);  // "light bulb" icon
          append(result.getText());
          return;
        }

        setIcon(AllIcons.Actions.EditSource);

        String baseName = ResultUtils.getBaseName(result.getCode().getPath());
        append(baseName, SimpleTextAttributes.REGULAR_ATTRIBUTES, true);

        int count = result.getCode().getLines().size();
        String format = " (%d result" + (count != 1 ? "s" : "") + ")";
        append(String.format(format, count), SUFFIX_ATTRIBUTES, false);
        return;
      }

      // This should never happen, but just in case IntelliJ gives us random nodes, then be sure
      // to always render something.
      append(value.toString());
    }
  }
}