package util.sqlparse.visitor.es.dense.data.query;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
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.tree.TerminalNode;
import org.antlr.v4.runtime.tree.TerminalNodeImpl;
import java.util.List;
import java.util.regex.Pattern;
import util.sqlparse.visitor.es.dense.Denseable;
import util.sqlparse.visitor.es.dense.EsBuilder;
import util.sqlparse.visitor.es.dense.ParamContext;
import util.sqlparse.visitor.es.memo.ApiResult;
import util.sqlparse.visitor.es.memo.FlowType;

public class SearchDense extends Denseable {
   @FlowType(FlowType.FlowTypeValue.down)
   public void denseTable(ParamContext.TableDenseContext params) {
      ApiResult api = params.getApi();
      EsParser.StatementContext statement = (EsParser.StatementContext)api.getStatement();
      EsParser.QueryDataStatementContext query = statement.queryDataStatement();
      if (query != null) {
         EsParser.QueryDataStatementSearchContext search = query.queryDataStatementSearch();
         if (search != null) {
            this.changeHttp(search, api);
            super.denseTable(params);
         }
      }
   }

   private <T extends EsParserContextBase> void changeHttp(EsParser.QueryDataStatementSearchContext ctx, ApiResult api) {
      EsParser.UriSegIndexContext index = ctx.indexName;
      EsParser.UriSegIndexContext mapping = ctx.mappingName;
      EsBuilder builder = EsBuilder.builder();
      EsParserContextBase uriIndex = this.buildEntityUri(builder, api.getSchemas());
      builder.reset();
      if (index != null) {
         builder.enter(index).replaceMe(uriIndex);
      } else {
         builder.enter(ctx).at(1).slash().append(2, uriIndex);
      }

      EsParserContextBase uriMapping = this.buildEntityUri(builder, api.getTables());
      builder.reset();
      if (mapping != null) {
         builder.enter(mapping).replaceMe(uriMapping);
      } else {
         builder.enter(ctx).at(3).slash().append(4, uriMapping);
      }

   }

   @FlowType(FlowType.FlowTypeValue.down)
   public void denseColumn(ParamContext.ColumnDenseContext params) {
      String replaceSchema = params.replaceSchema();
      String replaceTable = params.replaceTable();
      String replaceColumn = params.replaceColumn();
      String expression = params.expression();
      String response = params.response();
      Pattern pat = this.createColumnPattern(replaceSchema, replaceTable, replaceColumn);
      JSONObject data = JSON.parseObject(response);
      JSONObject dataHits = data.getJSONObject("hits");
      JSONArray hitsArray = dataHits.getJSONArray("hits");
      if (hitsArray != null && hitsArray.size() != 0) {
         for(Object hitsObject : hitsArray) {
            if (hitsObject instanceof JSONObject) {
               JSONObject hits = (JSONObject)hitsObject;
               String type = hits.getString("_type");
               String index = hits.getString("_index");
               JSONObject source = hits.getJSONObject("_source");
               StringBuilder path = new StringBuilder();
               path.append(index).append(".").append(type).append(".");
               this.denseColumnData(pat, source, expression, path, params);
            }
         }

         params.setResponse(data.toString());
      }
   }

   @FlowType(FlowType.FlowTypeValue.down)
   public void denseData(ParamContext.DenseAllContext params) {
      ParamContext.ColumnDenseContext columnDenseContext = new ParamContext.ColumnDenseContext(params);
      this.denseColumn(columnDenseContext);
      params.putAll(columnDenseContext);
   }

   @FlowType(FlowType.FlowTypeValue.up)
   public void denseRow(ParamContext.RowDenseContext params) {
      String operateType = params.operateType();
      if (operateType == "select") {
         ApiResult api = params.getApi();
         EsParser.StatementContext ctx = (EsParser.StatementContext)api.getStatement();
         String painless = this.buildPainless(params);
         if (painless != null && painless.length() != 0) {
            EsParser.QueryDataStatementContext query = ctx.queryDataStatement();
            EsParser.QueryDataStatementSearchContext search = query.queryDataStatementSearch();
            EsParser.QueryDataStatementSearchBodyContext searchBody = (EsParser.QueryDataStatementSearchBodyContext)EsBuilder.searchTree(search, EsParser.QueryDataStatementSearchBodyContext.class, 1, 5);
            this.maskQueryDataStatementSearchBodyQuery(painless, search, searchBody);
            this.limitQuery(search, params.limit());
         }
      }
   }

   private void limitQuery(EsParser.QueryDataStatementSearchContext search, String limit) {
      if (search != null && limit != null && limit.length() != 0) {
         EsParser.QueryDataStatementSearchBodyContext searchBody = (EsParser.QueryDataStatementSearchBodyContext)EsBuilder.searchTree(search, EsParser.QueryDataStatementSearchBodyContext.class, 1, 5);
         if (searchBody != null) {
            EsParser.StringContext ctxSize = null;
            EsParser.JvalueContext ctxValue = null;
            EsParser.PairContext keyValue = null;
            List<EsParser.PairsContext> pairs = searchBody.pairs();
            if (pairs != null && pairs.size() > 0) {
               for(EsParser.PairsContext ctxPairs : pairs) {
                  EsParser.PairContext pair = ctxPairs.pair();
                  if (pair != null) {
                     EsParser.StringContext key = pair.string();
                     if (key.getText().equalsIgnoreCase("\"size\"")) {
                        keyValue = pair;
                        ctxSize = pair.string();
                        ctxValue = pair.jvalue();
                     }
                  }
               }
            }

            EsBuilder builder = EsBuilder.builder();
            if (keyValue == null) {
               builder.reset();
               keyValue = (EsParser.PairContext)builder.enter(searchBody).at(1).pairs().enter().pair().enter().current();
            }

            if (ctxSize == null) {
               builder.reset();
               builder.enter(keyValue).rear(1).stringc("size");
            }

            if (ctxValue == null) {
               builder.reset();
               TerminalNodeImpl value = builder.terminal(278);
               builder.enter(keyValue).rear().jvalue().enter().append(value);
               CommonToken token = (CommonToken)value.getSymbol();
               token.setText(limit);
            } else {
               CommonToken token = null;
               TerminalNode number = ctxValue.NUMBER();
               if (number != null) {
                  token = (CommonToken)number.getSymbol();
               } else {
                  EsParser.StringContext str = ctxValue.string();
                  TerminalNode string = str.STRING();
                  token = (CommonToken)string.getSymbol();
               }

               if (token != null) {
                  String text = token.getText();
                  int cur = Integer.parseInt(text.trim());
                  int lim = Integer.parseInt(limit.trim());
                  int min = cur < lim ? cur : lim;
                  token.setText(min + "");
               }
            }

         }
      }
   }
}
