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.tree.ParseTree;
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;
import util.sqlparse.visitor.es.memo.IndexMemo;
import util.sqlparse.visitor.es.memo.MappingMemo;
import util.sqlparse.visitor.es.memo.Reference;

public class MSearchDense extends Denseable {
   @FlowType(FlowType.FlowTypeValue.up)
   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.QueryDataStatementMSearchContext msearch = query.queryDataStatementMSearch();
         if (msearch != null) {
            this.changeHttp(msearch, api);
            this.changeHeadParts(msearch, api);
            super.denseTable(params);
         }
      }
   }

   private void changeHttp(EsParser.QueryDataStatementMSearchContext 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();
      int k = 0;
      if (index != null) {
         builder.enter(index).replaceMe(uriIndex);
      } else {
         builder.enter(ctx).at(k + 1).slash().append(k + 2, uriIndex);
         k += 2;
      }

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

   }

   private void changeHeadParts(EsParser.QueryDataStatementMSearchContext ctx, ApiResult api) {
      EsParser.QueryDataStatementMSearchBodyContext body = ctx.queryDataStatementMSearchBody();
      if (body.children != null && body.children.size() != 0) {
         EsBuilder builder = EsBuilder.builder();
         List<IndexMemo> indices = api.getSchemas();
         List<MappingMemo> mappings = api.getTables();

         for(int i = 0; i < body.children.size(); i += 2) {
            ParseTree child = (ParseTree)body.children.get(i);
            builder.reset();
            if (child instanceof EsParser.QueryDataStatementMSearchEmptyContext) {
               builder.enter(body).at(i + 1).queryDataStatementSearchBody().enter().rear().queryDataStatementSearchBodyIndexTypes().enter().rear();
               this.buildIndexTypeClause(builder, indices, mappings);
               body.children.remove(i);
            } else if (child instanceof EsParser.QueryDataStatementSearchBodyContext) {
               EsParser.QueryDataStatementSearchBodyContext searchBody = (EsParser.QueryDataStatementSearchBodyContext)child;
               List<EsParser.QueryDataStatementSearchBodyIndexTypesContext> indexTypes = searchBody.queryDataStatementSearchBodyIndexTypes();
               if (indexTypes != null && indexTypes.size() != 0) {
                  EsParser.QueryDataStatementSearchBodyIndexTypesContext indexType = (EsParser.QueryDataStatementSearchBodyIndexTypesContext)indexTypes.get(0);
                  List<EsParser.QueryDataStatementSearchBodyIndexContext> indexCtxs = indexType.queryDataStatementSearchBodyIndex();
                  List<EsParser.QueryDataStatementSearchBodyTypeContext> typeCtxs = indexType.queryDataStatementSearchBodyType();
                  builder.enter(indexType);
                  if (indexCtxs == null || indexCtxs.size() == 0) {
                     builder.last().queryDataStatementSearchBodyIndex().enter().last();
                     this.buildEntityArray(builder, indices);
                  }

                  if (typeCtxs == null || typeCtxs.size() == 0) {
                     builder.last().queryDataStatementSearchBodyType().enter().last();
                     this.buildEntityArray(builder, mappings);
                  }
               } else {
                  builder.enter(searchBody);
                  this.buildIndexTypeClause(builder, indices, mappings);
               }
            }
         }

      }
   }

   private void buildIndexTypeClause(EsBuilder builder, List<IndexMemo> indices, List<MappingMemo> mappings) {
      builder.at(-1).queryDataStatementSearchBodyIndex().enter().at(-1);
      this.buildEntityArray(builder, indices);
      builder.pop();
      builder.at(-1).queryDataStatementSearchBodyType().enter().at(-1);
      this.buildEntityArray(builder, mappings);
      builder.pop();
   }

   private <T extends Reference> void buildEntityArray(EsBuilder builder, List<T> entities) {
      builder.stringArr();
      if (entities != null && entities.size() != 0) {
         builder.enter().rear();

         for(T entity : entities) {
            if (entity != null && entity.name != null) {
               EsParserContextBase ref = builder.stringc(entity.name).current();
               entity.refs.add(ref);
            }
         }

         builder.pop();
      }
   }

   @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);
      JSONArray responses = data.getJSONArray("responses");

      for(int i = 0; i < responses.size(); ++i) {
         JSONObject sets = responses.getJSONObject(i);
         JSONObject dataHits = sets.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();
            if (query != null) {
               EsParser.QueryDataStatementMSearchContext search = query.queryDataStatementMSearch();
               if (search != null) {
                  EsParser.QueryDataStatementMSearchBodyContext body = search.queryDataStatementMSearchBody();
                  if (body != null) {
                     EsBuilder builder = EsBuilder.builder();
                     if (body.children.size() % 2 != 0) {
                        builder.enter(body).queryDataStatementSearchBody();
                     }

                     builder.reset();

                     for(int i = 1; i < body.children.size(); i += 2) {
                        ParseTree child = (ParseTree)body.children.get(i);
                        if (child instanceof EsParser.QueryDataStatementMSearchEmptyContext) {
                           builder.enter(body).at(i + 1).queryDataStatementSearchBody();
                           body.children.remove(i);
                        }

                        EsParser.QueryDataStatementSearchBodyContext indexTypeSearchBody = (EsParser.QueryDataStatementSearchBodyContext)body.children.get(i);
                        this.maskQueryDataStatementSearchBodyQuery(painless, body, indexTypeSearchBody);
                     }

                  }
               }
            }
         }
      }
   }
}
