package com.karateca.ddescriber.dialog;

import com.intellij.openapi.project.Project;
import com.intellij.openapi.ui.DialogWrapper;
import com.intellij.ui.SpeedSearchComparator;
import com.intellij.ui.TreeSpeedSearch;
import com.intellij.ui.components.JBScrollPane;
import com.intellij.ui.treeStructure.Tree;
import com.karateca.ddescriber.Hierarchy;
import com.karateca.ddescriber.TestFindResult;
import org.jetbrains.annotations.Nullable;

import javax.swing.*;
import javax.swing.tree.DefaultMutableTreeNode;
import javax.swing.tree.TreePath;
import java.awt.event.*;
import java.util.ArrayList;
import java.util.Enumeration;
import java.util.List;
import java.util.Stack;

/**
 * @author Andres Dominguez.
 */
public class TreeViewDialog extends DialogWrapper {
  public static final int CLEAN_CURRENT_EXIT_CODE = 100;
  public static final int GO_TO_TEST_EXIT_CODE = 101;

  protected final Hierarchy hierarchy;
  protected static final int VISIBLE_ROW_COUNT = 13;
  private Tree tree;
  private TestFindResult selectedTest;

  public TreeViewDialog(@Nullable Project project, Hierarchy hierarchy) {
    super(project);
    this.hierarchy = hierarchy;
    init();
    setTitle("Select the Test or Suite to Add / Remove");
  }

  @Nullable
  @Override
  protected JComponent createCenterPanel() {
    List<TestFindResult> elements = hierarchy.getAllUnitTests();
    final TestFindResult closest = hierarchy.getClosest();

    DefaultMutableTreeNode root = populateTree(elements);

    tree = new Tree(root);

    tree.setVisibleRowCount(VISIBLE_ROW_COUNT);
    tree.setCellRenderer(new CustomTreeCellRenderer());

    // Add search, make it case insensitive.
    new TreeSpeedSearch(tree) {
      @Override
      protected boolean compare(String text, String pattern) {
        return super.compare(text.toLowerCase(), pattern.toLowerCase());
      }
    }.setComparator(new SpeedSearchComparator(false));

    addKeyAndMouseEvents();

    JBScrollPane scrollPane = new JBScrollPane(tree);
    selectClosestTest(root, closest);

    return scrollPane;
  }

  private void addKeyAndMouseEvents() {
    // Perform the OK action on enter.
    tree.addKeyListener(new KeyAdapter() {
      @Override
      public void keyPressed(KeyEvent keyEvent) {
        if (keyEvent.getKeyCode() == KeyEvent.VK_ENTER) {
          doOKAction();
        }
      }
    });

    // Go to the test on double click.
    tree.addMouseListener(new MouseAdapter() {
      @Override
      public void mousePressed(MouseEvent e) {
        int selRow = tree.getRowForLocation(e.getX(), e.getY());
        if (selRow != -1 && e.getClickCount() == 2) {
          TreePath selPath = tree.getPathForLocation(e.getX(), e.getY());
          nodeWasDoubleClicked(selPath);
        }
      }
    });
  }

  private void nodeWasDoubleClicked(TreePath selPath) {
    DefaultMutableTreeNode lastPathComponent = (DefaultMutableTreeNode) selPath.getLastPathComponent();
    selectedTest = (TestFindResult) lastPathComponent.getUserObject();
    close(GO_TO_TEST_EXIT_CODE);
  }

  private void selectClosestTest(DefaultMutableTreeNode root, final TestFindResult closest) {
    Enumeration enumeration = root.breadthFirstEnumeration();

    while (enumeration.hasMoreElements()) {
      DefaultMutableTreeNode node = (DefaultMutableTreeNode) enumeration.nextElement();

      if (node.getUserObject() == closest) {
        TreePath treePath = new TreePath(node.getPath());
        tree.setSelectionPath(treePath);
        tree.scrollPathToVisible(treePath);

        return;
      }
    }
  }

  private DefaultMutableTreeNode populateTree(List<TestFindResult> elements) {
    TestFindResult first = elements.get(0);
    DefaultMutableTreeNode root = new DefaultMutableTreeNode(first);

    if (elements.size() < 2) {
      return root;
    }

    Stack<DefaultMutableTreeNode> stack = new Stack<DefaultMutableTreeNode>();
    int currentIndentation = first.getIndentation();

    DefaultMutableTreeNode parent = root;
    DefaultMutableTreeNode last = root;

    for (TestFindResult element : elements.subList(1, elements.size())) {
      int ind = element.getIndentation();

      DefaultMutableTreeNode newNode = new DefaultMutableTreeNode(element);

      if (ind > currentIndentation) {
        stack.push(parent);
        parent = last;
      } else if (ind < currentIndentation) {
        parent = stack.pop();
      }
      last = newNode;
      parent.add(last);

      currentIndentation = ind;
    }

    return root;
  }

  @Nullable
  @Override
  public JComponent getPreferredFocusedComponent() {
    return tree;
  }

  public List<TestFindResult> getSelectedValues() {
    List<TestFindResult> selected = new ArrayList<TestFindResult>();

    for (DefaultMutableTreeNode node : tree.getSelectedNodes(DefaultMutableTreeNode.class, null)) {
      selected.add((TestFindResult) node.getUserObject());
    }

    return selected;
  }

  @Override
  protected Action[] createLeftSideActions() {
    return new Action[]{
        new DialogWrapperExitAction("Clean file", CLEAN_CURRENT_EXIT_CODE)
    };
  }

  public TestFindResult getSelectedTest() {
    return selectedTest;
  }
}

@@ -302,16 +302,59 @@
 e.Tree;%0A
+import com.karateca.ddescriber.ActionUtil;%0A
 import c
@@ -1496,16 +1496,27 @@
  root = 
+ActionUtil.
 populate
@@ -3606,948 +3606,8 @@
  %7D%0A%0A
-  private DefaultMutableTreeNode populateTree(List%3CTestFindResult%3E elements) %7B%0A    TestFindResult first = elements.get(0);%0A    DefaultMutableTreeNode root = new DefaultMutableTreeNode(first);%0A%0A    if (elements.size() %3C 2) %7B%0A      return root;%0A    %7D%0A%0A    Stack%3CDefaultMutableTreeNode%3E stack = new Stack%3CDefaultMutableTreeNode%3E();%0A    int currentIndentation = first.getIndentation();%0A%0A    DefaultMutableTreeNode parent = root;%0A    DefaultMutableTreeNode last = root;%0A%0A    for (TestFindResult element : elements.subList(1, elements.size())) %7B%0A      int ind = element.getIndentation();%0A%0A      DefaultMutableTreeNode newNode = new DefaultMutableTreeNode(element);%0A%0A      if (ind %3E currentIndentation) %7B%0A        stack.push(parent);%0A        parent = last;%0A      %7D else if (ind %3C currentIndentation) %7B%0A        parent = stack.pop();%0A      %7D%0A      last = newNode;%0A      parent.add(last);%0A%0A      currentIndentation = ind;%0A    %7D%0A%0A    return root;%0A  %7D%0A%0A
   @N

Use ActionUtil