package com.andresdominguez.ngsort;

import com.intellij.lang.javascript.psi.JSBlockStatement;
import com.intellij.lang.javascript.psi.JSFunctionExpression;
import com.intellij.openapi.actionSystem.AnAction;
import com.intellij.openapi.actionSystem.AnActionEvent;
import com.intellij.openapi.actionSystem.CommonDataKeys;
import com.intellij.openapi.actionSystem.PlatformDataKeys;
import com.intellij.openapi.command.CommandProcessor;
import com.intellij.openapi.editor.Document;
import com.intellij.openapi.editor.Editor;
import com.intellij.openapi.util.TextRange;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
import com.intellij.psi.util.PsiTreeUtil;

public class ToArrowFunctionAction extends AnAction {

  @Override
  public void actionPerformed(AnActionEvent e) {
    final Editor editor = e.getData(PlatformDataKeys.EDITOR);
    PsiFile psiFile = e.getData(CommonDataKeys.PSI_FILE);
    if (editor == null || psiFile == null) {
      return;
    }

    PsiElement element = psiFile.findElementAt(editor.getCaretModel().getOffset());
    final JSFunctionExpression functionExpr = PsiTreeUtil.getParentOfType(element, JSFunctionExpression.class);
    if (functionExpr == null) {
      return;
    }

    final JSBlockStatement fnBlock = PsiTreeUtil.findChildOfType(functionExpr, JSBlockStatement.class);
    if (fnBlock == null) {
      return;
    }

    final Document document = editor.getDocument();
    CommandProcessor.getInstance().executeCommand(getEventProject(e), new Runnable() {
      @Override
      public void run() {
        // Add =>
        TextRange fnBlockRange = fnBlock.getTextRange();
        document.replaceString(fnBlockRange.getStartOffset(), fnBlockRange.getEndOffset(),
            String.format("=> %s", fnBlock.getText()));

        // Migrate function() to ()
        int startOffset = functionExpr.getTextRange().getStartOffset();
        int endOffset = startOffset + functionExpr.getText().indexOf("(");
        document.replaceString(startOffset, endOffset, "");
      }
    }, "ng sort", null);

  }
}

@@ -146,24 +146,81 @@
 Expression;%0A
+import com.intellij.lang.javascript.psi.JSParameterList;%0A
 import com.i
@@ -1271,32 +1271,186 @@
  return;%0A    %7D%0A%0A
+    final JSParameterList paramList = PsiTreeUtil.findChildOfType(functionExpr, JSParameterList.class);%0A    if (paramList == null) %7B%0A      return;%0A    %7D%0A%0A
     final JSBloc
@@ -1577,32 +1577,155 @@
  return;%0A    %7D%0A%0A
+    final TextRange fnBlockRange = fnBlock.getTextRange();%0A    final TextRange paramListRange = paramList.getTextRange();%0A%0A
     final Docume
@@ -1911,65 +1911,8 @@
  =%3E%0A
-        TextRange fnBlockRange = fnBlock.getTextRange();%0A
     
@@ -2103,53 +2103,148 @@
     
-int startOffset = functionExpr.getTex
+boolean shouldDropParens = paramList.getParameters().length == 1;%0A        if (shouldDropParens) %7B%0A          int startOffset = paramLis
 tRange
-()
 .get
@@ -2262,24 +2262,26 @@
 ();%0A        
+ 
+ 
 int endOffse
@@ -2288,56 +2288,333 @@
 t = 
-startOffset + functionExpr.getText().indexOf(%22(%22
+paramListRange.getEndOffset();%0A          String replacement = paramList.getText().replace(%22(%22, %22%22).replace(%22)%22, %22%22);%0A          document.replaceString(startOffset, endOffset, replacement);%0A        %7D%0A%0A        int startOffset = functionExpr.getTextRange().getStartOffset();%0A        int endOffset = paramListRange.getStartOffset(
 );%0A 

Add support for 1 param list