package com.tang.intellij.lua.folding;

import com.intellij.lang.ASTNode;
import com.intellij.lang.folding.FoldingBuilderEx;
import com.intellij.lang.folding.FoldingDescriptor;
import com.intellij.lang.folding.NamedFoldingDescriptor;
import com.intellij.openapi.editor.Document;
import com.intellij.openapi.editor.FoldingGroup;
import com.intellij.openapi.util.TextRange;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiWhiteSpace;
import com.intellij.psi.impl.source.tree.ElementType;
import com.intellij.psi.util.PsiTreeUtil;
import com.tang.intellij.lua.psi.LuaBlock;
import com.tang.intellij.lua.psi.LuaPsiTreeUtil;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;

/**
 * Created by tangzx
 * Date : 2015/11/16.
 */
public class LuaFoldingBuilder extends FoldingBuilderEx {

    static final String HOLDER_TEXT = "...";

    @NotNull
    @Override
    public FoldingDescriptor[] buildFoldRegions(@NotNull PsiElement psiElement, @NotNull Document document, boolean b) {
        List<FoldingDescriptor> descriptors = new ArrayList<>();
        Collection<LuaBlock> luaFuncBodies = PsiTreeUtil.findChildrenOfType(psiElement, LuaBlock.class);
        luaFuncBodies.forEach(block -> {
            PsiElement prev = PsiTreeUtil.skipSiblingsBackward(block, PsiWhiteSpace.class);
            PsiElement next = PsiTreeUtil.skipSiblingsForward(block, PsiWhiteSpace.class);
            if (prev != null && next != null) {
                int l = prev.getTextOffset() + prev.getTextLength();
                int r = next.getTextOffset();

                TextRange range = new TextRange(l, r);
                if (range.getLength() > 0 && !addOneLineMethodFolding(descriptors, block, range, prev, next)) {
                    descriptors.add(new FoldingDescriptor(block, range) {
                        @Nullable
                        @Override
                        public String getPlaceholderText() {
                            return HOLDER_TEXT;
                        }
                    });
                }
            }
        });
        return descriptors.toArray(new FoldingDescriptor[descriptors.size()]);
    }

    private boolean addOneLineMethodFolding(List<FoldingDescriptor> descriptors, LuaBlock block, TextRange range, PsiElement prev, PsiElement next) {
        PsiElement[] children = block.getChildren();
        if (children.length == 0) return false;

        int validCount = 0;
        PsiElement first = null;
        for (PsiElement child : children) {
            if (!(child instanceof PsiWhiteSpace)) {
                first = child;
                validCount++;
                if (validCount > 1) return false;
            }
        }

        if (validCount == 1) {
            if (first.textContains('\n')) return false;

            FoldingGroup group = FoldingGroup.newGroup("one-liner");

            TextRange lRange = new TextRange(range.getStartOffset(), first.getTextOffset());
            descriptors.add(new NamedFoldingDescriptor(prev.getNode(), lRange, group, HOLDER_TEXT));

            TextRange rRange = new TextRange(first.getTextOffset() + first.getTextLength(), range.getEndOffset());
            descriptors.add(new NamedFoldingDescriptor(next.getNode(), rRange, group, HOLDER_TEXT));

            return true;
        }

        return false;
    }

    @Nullable
    @Override
    public String getPlaceholderText(@NotNull ASTNode astNode) {
        return null;
    }

    @Override
    public boolean isCollapsedByDefault(@NotNull ASTNode astNode) {
        return true;
    }
}

@@ -114,18 +114,16 @@
 gBuilder
-Ex
 ;%0Aimport
@@ -466,25 +466,14 @@
 psi.
-impl.source.
 tree.
+I
 Elem
@@ -602,27 +602,21 @@
 .psi.Lua
-PsiTreeUtil
+Types
 ;%0Aimport
@@ -727,37 +727,8 @@
 st;%0A
-import java.util.Collection;%0A
 impo
@@ -833,14 +833,17 @@
 der 
-extend
+implement
 s Fo
@@ -858,10 +858,8 @@
 lder
-Ex
  %7B%0A%0A
@@ -992,29 +992,20 @@
 ull 
-PsiElement psiElement
+ASTNode node
 , @N
@@ -1032,19 +1032,8 @@
 ment
-, boolean b
 ) %7B%0A
@@ -1109,145 +1109,540 @@
     
-C
+c
 ollect
-ion%3CLuaBlock%3E luaFuncBodies = PsiTreeUtil.findChildrenOfType(psiElement, LuaBlock.class);%0A        luaFuncBodies.forEach(block -%3E %7B
+DescriptorsRecursively(node, document, descriptors);%0A        return descriptors.toArray(new FoldingDescriptor%5Bdescriptors.size()%5D);%0A    %7D%0A%0A    private  void collectDescriptorsRecursively(@NotNull ASTNode node,%0A                                                @NotNull Document document,%0A                                                @NotNull List%3CFoldingDescriptor%3E descriptors) %7B%0A        IElementType type = node.getElementType();%0A        if (type == LuaTypes.BLOCK) %7B%0A            LuaBlock block = (LuaBlock) node.getPsi();
 %0A   
@@ -2223,18 +2223,18 @@
 , range)
- %7B
+);
 %0A       
@@ -2242,24 +2242,26 @@
         
+%7D%0A
         
 @Nullabl
@@ -2256,30 +2256,22 @@
     
-@Nullable%0A
     
+%7D%0A
         
     
@@ -2266,28 +2266,26 @@
         
-    
+%7D%0A
         
 @Overrid
@@ -2280,78 +2280,127 @@
     
-@Override%0A                        public String getPlaceholderText() %7B
+else if (type == LuaTypes.DOC_COMMENT) %7B%0A            descriptors.add(new FoldingDescriptor(node, node.getTextRange()));
 %0A   
@@ -2400,32 +2400,35 @@
 ));%0A        
+%7D%0A%0A
         
             
@@ -2419,212 +2419,137 @@
     
-            return HOLDER_TEXT;%0A                        %7D%0A                    %7D);%0A                %7D%0A            %7D%0A        %7D);%0A        return descriptors.toArray(new FoldingDescriptor%5Bdescriptors.size()%5D);
+for (ASTNode child : node.getChildren(null)) %7B%0A            collectDescriptorsRecursively(child, document, descriptors);%0A        %7D
 %0A   
@@ -3818,32 +3818,210 @@
 Node astNode) %7B%0A
+        IElementType type = astNode.getElementType();%0A        if (type == LuaTypes.BLOCK) return HOLDER_TEXT;%0A        else if (type == LuaTypes.DOC_COMMENT) return %22/** ... */%22;%0A
         return n
@@ -4125,26 +4125,27 @@
      return 
-tru
+fals
 e;%0A    %7D%0A%7D%0A

optimize folding builder