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

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.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.regex.Pattern;
import util.StringJoin;
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.MappingMemo;

public class UpdateDataDense extends Denseable {
   @FlowType(FlowType.FlowTypeValue.up)
   public void denseTable(ParamContext.TableDenseContext params) {
      ApiResult api = params.getApi();
      EsParser.StatementContext statement = (EsParser.StatementContext)api.getStatement();
      EsParser.InsertDataStatementContext insert = statement.insertDataStatement();
      if (insert == null) {
         EsParser.UpdateDataStatementContext update = statement.updateDataStatement();
         if (update != null) {
            EsParser.UpdateDataStatementByQueryContext byQuery = update.updateDataStatementByQuery();
            if (byQuery != null) {
               super.denseTable(params);
            } else {
               EsParser.UpdateDataStatementByIdContext byId = update.updateDataStatementById();
               if (byId != null) {
                  EsParser.UpdateDataStatementPart1Context form1 = byId.updateDataStatementPart1();
                  if (form1 != null) {
                     super.denseTable(params);
                     return;
                  }

                  List<MappingMemo> tables = api.getTables();
                  EsBuilder builder = EsBuilder.builder();
                  EsParser.UpdateDataStatementPart2Context form2 = byId.updateDataStatementPart2();
                  EsParserContextBase uriMapping = this.buildEntityUri(builder, tables);
                  builder.enter(form2).enter(form2.mappingName).replaceMe(uriMapping);
                  super.denseTable(params);
               }

            }
         }
      }
   }

   @FlowType(FlowType.FlowTypeValue.down)
   public void denseColumn(ParamContext.ColumnDenseContext params) {
   }

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

   @FlowType(FlowType.FlowTypeValue.up)
   public void denseRow(ParamContext.RowDenseContext params) {
      ApiResult api = params.getApi();
      EsParser.StatementContext statement = (EsParser.StatementContext)api.getStatement();
      EsParser.InsertDataStatementContext insert = statement.insertDataStatement();
      if (insert == null) {
         EsParser.UpdateDataStatementContext update = statement.updateDataStatement();
         if (update != null) {
            EsParser.UpdateDataStatementByQueryContext byQuery = update.updateDataStatementByQuery();
            if (byQuery != null) {
               this.denseRowByQuery(params, byQuery);
            } else {
               EsParser.UpdateDataStatementByIdContext byId = update.updateDataStatementById();
               if (byId != null) {
                  this.denseRowById(params, byId);
               }

            }
         }
      }
   }

   private void denseRowByQuery(ParamContext.RowDenseContext params, EsParser.UpdateDataStatementByQueryContext byQuery) {
      EsParser.UpdateDataStatementByQueryBodyContext updateBody = byQuery.updateDataStatementByQueryBody();
      List<EsParser.QueryBodyContext> queryBodys = updateBody.queryBody();
      if (queryBodys != null && queryBodys.size() != 0) {
         EsParser.QueryBodyContext queryBody = (EsParser.QueryBodyContext)queryBodys.get(0);
         String operateType = params.operateType();
         if (operateType == "select") {
            String painless = this.buildPainless(params);
            if (painless != null && painless.length() != 0) {
               EsBuilder builder = EsBuilder.builder();
               this.maskQueryBody(painless, builder, queryBody);
            }
         }
      }
   }

   private void denseRowById(ParamContext.RowDenseContext params, EsParser.UpdateDataStatementByIdContext byId) {
      if (byId != null) {
         EsParser.UpdateDataStatementPart1Context noMappingCtx = byId.updateDataStatementPart1();
         EsParser.CuDataStatementBodyContext body = null;
         if (noMappingCtx != null) {
            body = noMappingCtx.cuDataStatementBody();
         }

         EsParser.UpdateDataStatementPart2Context withMappingCtx = byId.updateDataStatementPart2();
         if (withMappingCtx != null) {
            EsParser.UriSegIndexContext mappingName = withMappingCtx.mappingName;
            EsBuilder builder = EsBuilder.builder();
            EsParserContextBase update = builder.uriSegIndex().enter().uriIndexAtom("_update").current();
            builder.enter(mappingName).replaceMe(update);
            body = withMappingCtx.cuDataStatementBody();
         }

         if (body != null) {
            this.denseRowById(params, body);
         }
      }
   }

   private void denseRowById(ParamContext.RowDenseContext params, EsParser.CuDataStatementBodyContext body) {
      Map<String, List<String>> modifyTable = params.modifyTable();
      if (modifyTable != null && modifyTable.size() != 0) {
         Map<String, String> tableKeyWord = params.tableKeyWord();
         if (tableKeyWord != null && tableKeyWord.size() != 0) {
            String table = (String)tableKeyWord.keySet().iterator().next();
            String value = (String)((List)modifyTable.values().iterator().next()).get(0);
            String col = (String)tableKeyWord.values().iterator().next();
            Pattern pat = this.createColumnPattern(col);
            List<ColInfo> cols = new ArrayList();
            if (this.visit(pat, body.cuDataStatementValue(), cols, new ArrayList())) {
               if (cols.size() != 0) {
                  String script = this.buildUpdateDataPainless(table, value, cols);
                  if (script != null) {
                     EsBuilder builder = EsBuilder.builder();
                     builder.enter(body);
                     builder.remove(body.cuDataStatementValue());
                     builder.at(0).cuDataStatementObjValue().enter().rear().cuDataStatementObjValuePair().enter().cuDataStatementFieldScript(script);
                  }
               }
            }
         }
      }
   }

   String buildUpdateDataPainless(String table, String value, List<ColInfo> cols) {
      if (cols.size() == 0) {
         return null;
      } else {
         StringBuilder script = new StringBuilder();
         List<ColInfo> matched = new ArrayList();

         for(ColInfo col : cols) {
            if (col.matches) {
               matched.add(col);
            }
         }

         StringBuilder colLimitedScripts = new StringBuilder();

         for(ColInfo col : matched) {
            String fullName = StringJoin.join(col.fullNames, ".");
            colLimitedScripts.append(String.format("&& ctx._source['%s']!='%s'", fullName, value));
         }

         script.append(String.format("if(ctx._type ==~ /%s/ %s){\n", table, colLimitedScripts));
         int i = 0;

         for(ColInfo col : cols) {
            String sourceVal = col.value;
            if (col.fullNames.size() > 1) {
               for(int i1 = col.fullNames.size() - 1; i1 >= 1; --i1) {
                  String name = (String)col.fullNames.get(i1);
                  script.append(String.format("\tdef map%d=new HashMap();\n", i));
                  if (i1 == col.fullNames.size() - 1) {
                     script.append(String.format("\tmap%d.put('%s','%s');\n", i, name, col.value));
                  } else {
                     script.append(String.format("\tmap%d.put('%s',map%d);\n", i, name, i - 1));
                  }

                  ++i;
               }

               sourceVal = String.format("map%d", i - 1);
            }

            if (col.matches) {
               script.append("\t");
            }

            script.append(String.format("ctx._source['%s']=%s;\n", col.fullNames.get(0), sourceVal));
         }

         script.append("}\n");
         return script.toString();
      }
   }

   private boolean visit(Pattern pat, EsParserContextBase ctx, List<ColInfo> cols, List<String> path) {
      if (ctx.children != null && ctx.children.size() != 0) {
         for(ParseTree child : ctx.children) {
            if (child.getClass().equals(EsParser.CuDataStatementObjValuePairContext.class)) {
               EsParser.CuDataStatementObjValuePairContext pair = (EsParser.CuDataStatementObjValuePairContext)child;
               EsParser.CuDataStatementFieldScriptContext script = pair.script;
               if (script != null) {
                  cols.clear();
                  return false;
               }

               EsParser.StringContext fieldName = pair.fieldName;
               String field = fieldName.getTextValue();
               EsParser.CuDataStatementValueContext fieldValue = pair.fieldValue;
               List<String> childPath = new ArrayList(path);
               childPath.add(field);
               if (fieldValue.objVal == null) {
                  String text = fieldValue.getTextValue();
                  ColInfo col = new ColInfo();
                  col.name = field;
                  col.value = text;
                  col.fullNames = childPath;
                  String name = StringJoin.join(childPath, ".");
                  col.matches = pat.matcher(name).matches();
                  cols.add(col);
               } else if (!this.visit(pat, fieldValue.objVal, cols, childPath)) {
                  return false;
               }
            } else if (child instanceof EsParserContextBase && !this.visit(pat, (EsParserContextBase)child, cols, path)) {
               return false;
            }
         }

         return true;
      } else {
         return true;
      }
   }

   static class ColInfo {
      List<String> fullNames = new ArrayList();
      String name;
      String value;
      boolean matches;
   }
}
