package util.sqlparse.visitor.es.dense;

import com.alibaba.druid.sql.dialect.es.ast.EsParser;
import com.alibaba.druid.sql.dialect.es.ast.EsParserContextBase;
import org.antlr.v4.runtime.CommonToken;
import org.antlr.v4.runtime.ParserRuleContext;
import org.antlr.v4.runtime.RuleContext;
import org.antlr.v4.runtime.tree.ParseTree;
import org.antlr.v4.runtime.tree.TerminalNode;
import org.antlr.v4.runtime.tree.TerminalNodeImpl;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Stack;
import java.util.concurrent.Callable;

public class EsBuilder {
   private State parent;
   private State current;
   private Stack<State> stack = new Stack();

   private EsBuilder() {
   }

   public static EsBuilder builder() {
      return new EsBuilder();
   }

   public static EsBuilder builder(EsParserContextBase parent) {
      EsBuilder builder = new EsBuilder();
      State state = new State(parent);
      builder.stack.push(state);
      builder.parent = state;
      return builder;
   }

   public EsBuilder reset() {
      this.parent = null;
      this.stack.clear();
      this.current = null;
      return this;
   }

   public EsParserContextBase build() {
      return this.parent.ctx;
   }

   public <T extends EsParserContextBase> T current() {
      return (T)this.current.ctx;
   }

   public EsBuilder remove(EsParserContextBase ctx) {
      if (this.parent != null) {
         this.parent.ctx.children.remove(ctx);
      }

      ctx.setParent((RuleContext)null);
      return this;
   }

   public EsBuilder removeAll(Collection<EsParserContextBase> ctxs) {
      if (this.parent != null) {
         this.parent.ctx.children.removeAll(ctxs);
      }

      for(EsParserContextBase ctx : ctxs) {
         ctx.setParent((RuleContext)null);
      }

      return this;
   }

   public EsBuilder replaceMe(EsParserContextBase ctx, Callable<Boolean>... cbs) {
      if (this.parent == null) {
         return this;
      } else {
         ParserRuleContext me = this.parent.ctx;
         ParserRuleContext father = me.getParent();
         boolean changed = false;
         int index = father.children.indexOf(me);
         if (index != -1) {
            String before = me.toString();
            String current = ctx.toString();
            changed = !before.equals(current);
            father.children.remove(index);
            father.children.add(index, ctx);
         } else {
            changed = true;
            father.children.add(ctx);
         }

         if (changed && cbs != null && cbs.length > 0) {
            for(Callable<Boolean> cb : cbs) {
               try {
                  cb.call();
               } catch (Exception exception) {
                  exception.printStackTrace();
               }
            }
         }

         return this;
      }
   }

   public EsBuilder setString(String value, int type) {
      this.parent.ctx.children.clear();
      TerminalNode val = type == 0 ? this.cstring(value) : this.ctext(value);
      this.addChild(val);
      return this;
   }

   public EsBuilder queryBody() {
      EsParser.QueryBodyContext ctx = new EsParser.QueryBodyContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.addChild(this.terminal(90, "\"query\""));
      this.colon();
      this.lparen();
      this.rparen();
      this.pop();
      this.current(ctx);
      return this;
   }

   private void current(EsParserContextBase ctx) {
      this.current = new State(ctx);
   }

   public EsBuilder queryFilter() {
      EsParser.QueryFilterContext ctx = new EsParser.QueryFilterContext(this.getParent(), 0);
      this.addChild(ctx);
      this.current(ctx);
      return this;
   }

   public EsBuilder comboundFilter() {
      EsParser.ComboundFilterContext ctx = new EsParser.ComboundFilterContext(this.getParent(), 0);
      this.addChild(ctx);
      this.current(ctx);
      return this;
   }

   public EsBuilder queryDataStatementSearchBody(EsParserContextBase... children) {
      EsParser.QueryDataStatementSearchBodyContext ctx = new EsParser.QueryDataStatementSearchBodyContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.lparen();
      if (children != null && children.length == 0) {
         for(EsParserContextBase child : children) {
            this.addChild(child);
         }
      }

      this.rparen();
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder queryDataStatementSearchBodyIndexTypes() {
      EsParser.QueryDataStatementSearchBodyIndexTypesContext ctx = new EsParser.QueryDataStatementSearchBodyIndexTypesContext(this.getParent(), 0);
      this.addChild(ctx);
      this.current(ctx);
      return this;
   }

   public EsBuilder queryDataStatementSearchBodyIndex() {
      EsParser.QueryDataStatementSearchBodyIndexContext ctx = new EsParser.QueryDataStatementSearchBodyIndexContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.addChild(this.terminal(101, "\"index\""));
      this.colon();
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder queryDataStatementSearchBodyType() {
      EsParser.QueryDataStatementSearchBodyTypeContext ctx = new EsParser.QueryDataStatementSearchBodyTypeContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.addChild(this.terminal(106, "\"type\""));
      this.colon();
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder queryDataStatementSearchBodyQuery(EsParserContextBase... children) {
      EsParser.QueryDataStatementSearchBodyQueryContext ctx = new EsParser.QueryDataStatementSearchBodyQueryContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      if (children != null && children.length == 0) {
         for(EsParserContextBase child : children) {
            this.addChild(child);
         }
      }

      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder scriptFilter(String sc) {
      EsParser.ScriptFilterContext ctx = new EsParser.ScriptFilterContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      TerminalNode script = this.terminal(86, "\"script\"");
      this.addChild(script);
      this.colon();
      this.lparen();
      this.addChild(script);
      this.colon();
      this.scriptFilterSourceValue(sc);
      this.rparen();
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder scriptFilterSourceValue(String script) {
      EsParser.ScriptFilterSourceValueContext ctx = new EsParser.ScriptFilterSourceValueContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.lparen();
      this.addChild(this.terminal(104, "\"source\""));
      this.colon();
      this.scriptc(script);
      this.rparen();
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder boolFilter() {
      EsParser.BoolFilterContext ctx = new EsParser.BoolFilterContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.addChild(this.terminal(165, "\"bool\""));
      this.colon();
      this.lparen();
      this.rparen();
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder mustFilter(boolean not) {
      EsParser.MustFilterContext ctx = new EsParser.MustFilterContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.addChild(this.terminal(not ? 95 : 95, not ? "\"must_not\"" : "\"must\""));
      this.colon();
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder mustFilter1() {
      EsParser.MustFilter1Context ctx = new EsParser.MustFilter1Context(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.lparen();
      this.rparen();
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder filterFilter() {
      EsParser.FilterFilterContext ctx = new EsParser.FilterFilterContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.addChild(this.terminal(98, "\"filter\""));
      this.colon();
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder filterFilter1() {
      EsParser.FilterFilter1Context ctx = new EsParser.FilterFilter1Context(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.lparen();
      this.rparen();
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder stringArr() {
      EsParser.StringArrContext ctx = new EsParser.StringArrContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.clsquare();
      this.crsquare();
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder stringArr(List<String> values) {
      EsParser.StringArrContext ctx = new EsParser.StringArrContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.clsquare();
      if (values != null && values.size() > 0) {
         for(String value : values) {
            EsParser.StringContext c = new EsParser.StringContext(ctx, 0);
            this.addChild(c);
            this.enter(c);
            this.addChild(this.ctext(value));
            this.pop();
         }
      }

      this.crsquare();
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder stringc(String value) {
      EsParser.StringContext c = new EsParser.StringContext(this.getParent(), 0);
      this.addChild(c);
      this.enter(c);
      this.addChild(this.cstring(value));
      this.pop();
      this.current(c);
      return this;
   }

   public EsBuilder textc(String value) {
      EsParser.StringContext c = new EsParser.StringContext(this.getParent(), 0);
      this.addChild(c);
      this.enter(c);
      this.addChild(this.ctext(value));
      this.pop();
      this.current(c);
      return this;
   }

   public EsBuilder scriptc(String value) {
      EsParser.StringContext c = new EsParser.StringContext(this.getParent(), 0);
      this.addChild(c);
      this.enter(c);
      this.addChild(this.cstring(value));
      this.pop();
      this.current(c);
      return this;
   }

   private ParserRuleContext getParent() {
      return this.parent == null ? null : this.parent.ctx;
   }

   public EsBuilder append(int pos, EsParserContextBase... ctxs) {
      if (ctxs != null && ctxs.length != 0) {
         if (this.parent.ctx.children == null) {
            this.parent.ctx.children = new ArrayList();
         }

         for(int i = 0; i < ctxs.length; ++i) {
            EsParserContextBase ctx = ctxs[i];
            if (ctx != null) {
               ctx.setParent(this.parent.ctx);
               this.parent.ctx.children.add(i + pos, ctx);
            }
         }

         return this;
      } else {
         return this;
      }
   }

   public EsBuilder append(EsParserContextBase... ctxs) {
      if (ctxs != null && ctxs.length != 0) {
         if (this.parent.ctx.children == null) {
            this.parent.ctx.children = new ArrayList();
         }

         for(int i = 0; i < ctxs.length; ++i) {
            EsParserContextBase ctx = ctxs[i];
            ctx.setParent(this.parent.ctx);
            this.parent.ctx.children.add(ctx);
         }

         return this;
      } else {
         return this;
      }
   }

   public EsBuilder append(EsParserContextBase ctx, int pos) {
      ctx.setParent(this.getParent());
      if (this.parent.ctx.children == null) {
         this.parent.ctx.children = new ArrayList();
      }

      this.parent.ctx.children.add(pos, ctx);
      return this;
   }

   public EsBuilder append(EsParserContextBase ctx, Class<?> from) {
      if (from == null) {
         this.parent.ctx.children.add(ctx);
      } else {
         for(int i = 0; i < this.parent.ctx.children.size(); ++i) {
            ParseTree c = (ParseTree)this.parent.ctx.children.get(i);
            if (c.getClass().equals(from)) {
               this.parent.ctx.children.add(i + 1, c);
               break;
            }
         }
      }

      return this;
   }

   public EsBuilder append(TerminalNode... nodes) {
      if (nodes != null && nodes.length != 0) {
         if (this.parent.ctx.children == null) {
            this.parent.ctx.children = new ArrayList();
         }

         for(int i = 0; i < nodes.length; ++i) {
            TerminalNode ctx = nodes[i];
            ctx.setParent(this.parent.ctx);
            this.parent.ctx.children.add(ctx);
         }

         return this;
      } else {
         return this;
      }
   }

   public EsBuilder pairs() {
      EsParser.PairsContext ctx = new EsParser.PairsContext(this.getParent(), 0);
      this.addChild(ctx);
      this.current(ctx);
      return this;
   }

   public EsBuilder pair() {
      EsParser.PairContext ctx = new EsParser.PairContext(this.getParent(), 0);
      this.addChild(ctx);
      this.current(ctx);
      return this;
   }

   public EsBuilder jvalue() {
      EsParser.JvalueContext ctx = new EsParser.JvalueContext(this.getParent(), 0);
      this.addChild(ctx);
      this.current(ctx);
      return this;
   }

   public EsBuilder uriSeg() {
      EsParser.UriSegContext ctx = new EsParser.UriSegContext(this.getParent(), 0);
      this.addChild(ctx);
      this.current(ctx);
      return this;
   }

   public EsBuilder uriSeg(String value) {
      EsParser.UriSegContext ctx = new EsParser.UriSegContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.addChild(this.terminal(279, value));
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder uriSegIndex() {
      EsParser.UriSegIndexContext ctx = new EsParser.UriSegIndexContext(this.getParent(), 0);
      this.addChild(ctx);
      this.current(ctx);
      return this;
   }

   public EsBuilder uriSegIndex(String value) {
      EsParser.UriSegIndexContext ctx = new EsParser.UriSegIndexContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.uriIndexAtom(value);
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder uriIndexAtom(String value) {
      EsParser.UriIndexAtomContext ctx = new EsParser.UriIndexAtomContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.addChild(this.terminal(279, value));
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder lparen() {
      TerminalNode node = this.clparen();
      this.addChild(node);
      return this;
   }

   public EsBuilder rparen() {
      TerminalNode node = this.crparen();
      this.addChild(node);
      return this;
   }

   public EsBuilder lsquare() {
      TerminalNode node = this.clsquare();
      this.addChild(node);
      return this;
   }

   public EsBuilder rsquare() {
      TerminalNode node = this.crsquare();
      this.addChild(node);
      return this;
   }

   public EsBuilder colon() {
      TerminalNode node = this.ccolon();
      this.addChild(node);
      return this;
   }

   public EsBuilder slash() {
      TerminalNode node = this.cslash();
      this.addChild(node);
      return this;
   }

   public EsBuilder leafFilter() {
      EsParser.LeafFilterContext ctx = new EsParser.LeafFilterContext(this.getParent(), 0);
      this.addChild(ctx);
      this.current(ctx);
      return this;
   }

   public EsBuilder leafOtherFilter() {
      EsParser.LeafFilterContext ctx = new EsParser.LeafFilterContext(this.getParent(), 0);
      this.addChild(ctx);
      this.current(ctx);
      return this;
   }

   public EsBuilder query() {
      TerminalNode query = this.terminal(90, "\"query\"");
      this.addChild(query);
      return this;
   }

   public EsBuilder cuDataStatementValue() {
      EsParser.CuDataStatementValueContext ctx = new EsParser.CuDataStatementValueContext(this.getParent(), 0);
      this.addChild(ctx);
      this.current(ctx);
      return this;
   }

   public EsBuilder cuDataStatementObjValue() {
      EsParser.CuDataStatementObjValueContext ctx = new EsParser.CuDataStatementObjValueContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.lparen();
      this.rparen();
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder cuDataStatementArrayValue() {
      EsParser.CuDataStatementObjValueContext ctx = new EsParser.CuDataStatementObjValueContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.lsquare();
      this.rsquare();
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder cuDataStatementObjValuePair() {
      EsParser.CuDataStatementObjValuePairContext ctx = new EsParser.CuDataStatementObjValuePairContext(this.getParent(), 0);
      this.addChild(ctx);
      this.current(ctx);
      return this;
   }

   public EsBuilder cuDataStatementFieldScript(String script) {
      EsParser.CuDataStatementFieldScriptContext ctx = new EsParser.CuDataStatementFieldScriptContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.addChild(this.terminal(86, "\"script\""));
      this.colon();
      this.scriptc(script);
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder queryDataStatementCountBody() {
      EsParser.QueryDataStatementCountBodyContext ctx = new EsParser.QueryDataStatementCountBodyContext(this.getParent(), 0);
      this.addChild(ctx);
      this.enter(ctx);
      this.lparen();
      this.rparen();
      this.pop();
      this.current(ctx);
      return this;
   }

   public EsBuilder script() {
      EsParser.CuDataStatementValueContext ctx = new EsParser.CuDataStatementValueContext(this.getParent(), 0);
      this.addChild(ctx);
      this.current(ctx);
      return this;
   }

   public EsBuilder term(int type, String value) {
      TerminalNode query = this.terminal(type, String.format("\"%s\"", value));
      this.addChild(query);
      return this;
   }

   public TerminalNodeImpl terminal(int token) {
      return new TerminalNodeImpl(new CommonToken(token, EsParser.ruleNames[token]));
   }

   private TerminalNode ccolon() {
      return this.terminal(12, ":");
   }

   private TerminalNode cslash() {
      return this.terminal(7, "/");
   }

   private TerminalNode clsquare() {
      return this.terminal(10, "[");
   }

   private TerminalNode crsquare() {
      return this.terminal(11, "]");
   }

   private TerminalNode clparen() {
      return this.terminal(8, "{");
   }

   private TerminalNode crparen() {
      return this.terminal(9, "}");
   }

   private TerminalNode cstring(String value) {
      return this.terminal(277, String.format("\"%s\"", value));
   }

   private TerminalNode cscript(String value) {
      return this.terminal(277, String.format("\"\"\"%s\"\"\"", value));
   }

   private TerminalNode ctext(String value) {
      return this.terminal(277, String.format("%s", value));
   }

   private TerminalNode terminal(int token, String value) {
      return new TerminalNodeImpl(new CommonToken(token, value));
   }

   private void addChild(ParseTree ctx) {
      if (this.parent != null) {
         EsParserContextBase c = this.parent.ctx;
         if (c.children == null) {
            c.children = new ArrayList();
         }

         int pos = Math.max(0, this.parent.mark < 0 ? c.children.size() : this.parent.mark);
         c.children.add(pos, ctx);
         ctx.setParent(c);
         if (this.parent.mark > -1) {
            ++this.parent.mark;
         }
      }

   }

   public EsBuilder enter() {
      if (this.current == null) {
         return this;
      } else {
         this.enter(this.current.ctx);
         return this;
      }
   }

   public EsBuilder enter(int mark, EsParserContextBase data) {
      if (this.parent == null) {
         this.parent = new State(data, mark);
         data.putAttribute("stack", 1);
         this.enter(data);
         return this;
      } else {
         if (this.stack.size() > 0) {
            State state = (State)this.stack.peek();
            EsParserContextBase c = state.ctx;
            if (c.equals(data)) {
               return this;
            }
         }

         State state = new State(data, mark);
         this.stack.push(state);
         this.parent = state;
         return this;
      }
   }

   public EsBuilder enter(EsParserContextBase data) {
      if (this.parent == null) {
         this.parent = new State(data);
         data.putAttribute("stack", 1);
         this.enter(data);
         return this;
      } else {
         if (this.stack.size() > 0) {
            State state = (State)this.stack.peek();
            EsParserContextBase c = state.ctx;
            if (c.equals(data)) {
               return this;
            }
         }

         State state = new State(data);
         this.stack.push(state);
         this.parent = state;
         return this;
      }
   }

   public EsBuilder at(int mark) {
      this.parent.mark = mark;
      return this;
   }

   public EsBuilder rear() {
      if (this.parent != null && this.parent.ctx != null) {
         if (this.parent.ctx.children == null) {
            this.parent.ctx.children = new ArrayList();
         }

         this.parent.mark = Math.max(0, this.parent.ctx.children.size() - 1);
         return this;
      } else {
         return this;
      }
   }

   public EsBuilder last() {
      if (this.parent != null && this.parent.ctx != null) {
         if (this.parent.ctx.children == null) {
            this.parent.ctx.children = new ArrayList();
         }

         this.parent.mark = Math.max(0, this.parent.ctx.children.size());
         return this;
      } else {
         return this;
      }
   }

   public EsBuilder rear(int minus) {
      if (this.parent != null && this.parent.ctx != null) {
         if (this.parent.ctx.children == null) {
            this.parent.ctx.children = new ArrayList();
         }

         this.parent.mark = Math.max(0, this.parent.ctx.children.size() - minus);
         return this;
      } else {
         return this;
      }
   }

   public EsBuilder pop() {
      if (this.stack.size() == 0) {
         return this;
      } else {
         State state = (State)this.stack.peek();
         EsParserContextBase c = state.ctx;
         if (c.getAttribute("stack") != null) {
            return this;
         } else {
            this.stack.pop();
            this.parent = (State)this.stack.peek();
            return this;
         }
      }
   }

   public static <T extends EsParserContextBase> T searchTree(EsParserContextBase ctx, Class<T> cls, int depth, int max) {
      if (ctx == null) {
         return null;
      } else if (ctx.children != null && ctx.children.size() != 0) {
         if (depth > max) {
            return null;
         } else {
            for(ParseTree child : ctx.children) {
               if (child.getClass().equals(cls)) {
                  return (T)(child);
               }

               if (child instanceof EsParserContextBase) {
                  EsParserContextBase find = searchTree((EsParserContextBase)child, cls, depth + 1, max);
                  if (find != null) {
                     return (T)find;
                  }
               }
            }

            return null;
         }
      } else {
         return null;
      }
   }

   static class State {
      int mark;
      EsParserContextBase ctx;

      State(EsParserContextBase ctx) {
         this(ctx, -1);
      }

      State(EsParserContextBase ctx, int mark) {
         this.mark = -1;
         this.ctx = ctx;
         this.mark = mark;
      }
   }
}
