content: /*
   * Copyright (C) 2016-2023 ActionTech.
   * License: http://www.gnu.org/licenses/gpl.html GPL version 2 or higher.
   */
  
  package com.actiontech.dble.backend.mysql.nio.handler.builder.sqlvisitor;
  
  import com.actiontech.dble.plan.common.item.Item;
  import com.actiontech.dble.plan.common.item.ItemBasicConstant;
  import com.actiontech.dble.plan.common.item.ItemNull;
  import com.actiontech.dble.plan.common.item.ItemString;
  import com.actiontech.dble.plan.common.item.function.ItemFunc;
  import com.actiontech.dble.plan.common.item.function.operator.ItemBoolFunc2;
  import com.actiontech.dble.plan.common.item.function.operator.cmpfunc.ItemFuncEqual;
  import com.actiontech.dble.plan.common.item.function.operator.cmpfunc.ItemFuncIn;
  import com.actiontech.dble.plan.common.item.function.operator.cmpfunc.ItemFuncNe;
  import com.actiontech.dble.plan.common.item.function.operator.logic.ItemCondAnd;
  import com.actiontech.dble.plan.common.item.function.operator.logic.ItemCondOr;
  import com.actiontech.dble.plan.common.item.function.operator.logic.ItemFuncNot;
  import com.actiontech.dble.plan.common.ptr.BoolPtr;
  import com.actiontech.dble.plan.node.ModifyNode;
  import com.actiontech.dble.plan.node.PlanNode;
  import com.actiontech.dble.plan.node.TableNode;
  import com.actiontech.dble.plan.util.PlanUtil;
  import com.actiontech.dble.route.RouteResultset;
  import com.actiontech.dble.server.parser.ServerParse;
  import com.actiontech.dble.util.StringUtil;
  import com.alibaba.druid.sql.ast.expr.SQLBinaryOperator;
  import org.jetbrains.annotations.NotNull;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.List;
  
  import static com.actiontech.dble.backend.mysql.nio.handler.query.impl.subquery.UpdateSubQueryHandler.NEED_REPLACE;
  
  
  public class UpdateVisitor extends MysqlVisitor {
  
      private static final Logger LOGGER = LoggerFactory.getLogger(UpdateVisitor.class);
      private final List<Item> valueItemList;
      private final List<Item> fieldList;
      protected boolean isExplain;
  
      public UpdateVisitor(PlanNode update, boolean isTopQuery, List<Item> valueItemList, List<Item> fieldList, boolean isExplain) {
          super(update, isTopQuery);
          this.valueItemList = valueItemList;
          this.fieldList = fieldList;
          this.isExplain = isExplain;
      }
  
      public void visit() {
          if (!visited) {
              replaceableSqlBuilder.clear();
              sqlBuilder = replaceableSqlBuilder.getCurrentElement().getSb();
              // if visited,push down visitor need just replace the name
              PlanNode.PlanNodeType i = query.type();
              if (i == PlanNode.PlanNodeType.MODIFY) {
                  visit((ModifyNode) query);
  
              } else {
                  throw new RuntimeException("not implement yet!");
              }
              visited = true;
          }
      }
  
      protected void visit(ModifyNode update) {
          sqlBuilder.append("update");
  
          //from
          update.getReferedTableNodes().stream()
                  .forEach(tableNode -> {
                      if (update.getSetItemList().stream()
                              .anyMatch(setItem -> !StringUtil.isEmpty(setItem.arguments().get(0).getTableName()) && (setItem.arguments().get(0).getTableName().equals(tableNode.getAlias()) || setItem.arguments().get(0).getTableName().equals(tableNode.getTableName())))) {
                          buildTableName(tableNode, sqlBuilder);
                      }
                  });
  
          //set
          sqlBuilder.append(" set ");
          String tableName = null;
          StringBuilder setBuilder = new StringBuilder();
          for (ItemFuncEqual itemFuncEqual : update.getSetItemList()) {
              Item setItem = itemFuncEqual.arguments().get(0);
              tableName = setItem.getTableName();
              if (StringUtil.isEmpty(tableName)) {
                  setBuilder.append("`" + setItem.getItemName() + "`");
              } else {
                  setBuilder.append("`" + setItem.getTableName() + "`.`" + setItem.getItemName() + "`");
              }
              setBuilder.append(" = ");
              Item valueItem = itemFuncEqual.arguments().get(1);
              String sqlStr = buildSQLStr(valueItem, tableName);
              setBuilder.append(sqlStr);
              setBuilder.append(",");
          }
          setBuilder.deleteCharAt(setBuilder.length() - 1);
          sqlBuilder.append(setBuilder);
  
          //where
          Item whereFilter = update.getWhereFilter();
          if (whereFilter != null) {
              Item item = rebuildUpdateItem(whereFilter, tableName);
              String selName = getUpdateItemName(item);
              sqlBuilder.append(" where ").append(selName);
          }
      }
  
      private String buildSQLStr(Item valueItem, String tableName) {
          if (valueItem instanceof ItemBasicConstant) {
              return valueItem.toString();
          } else if (!StringUtil.equalsIgnoreCase(tableName, valueItem.getTableName()) && isExplain) {
              return new ItemString(NEED_REPLACE, valueItem.getCharsetIndex()).toString();
          } else if (StringUtil.equalsIgnoreCase(tableName, valueItem.getTableName())) {
              return valueItem.toExpression().toString();
          } else {
              if (valueItemList.size() == 1) {
                  //autoalias_scalar
                  return valueItemList.get(0).toString();
              }
              int index = getItemIndex(valueItem);
              if (index < 0) {
                  return valueItem.toExpression().toString();
              } else {
                  return valueItemList.get(index).toString();
              }
          }
      }
  
      @Override
      protected String visitPushDownNameSel(Item o) {
          return null;
      }
  
  
      // try to trim sharding from field_item
      protected String getUpdateItemName(Item item) {
          if (item instanceof ItemCondOr) {
              StringBuilder sb = new StringBuilder();
              sb.append("(");
              for (int index = 0; index < item.getArgCount(); index++) {
                  if (index > 0) {
                      sb.append(" OR ");
                  }
                  sb.append(getUpdateItemName(item.arguments().get(index)));
              }
              sb.append(")");
              return sb.toString();
          } else if (item instanceof ItemCondAnd) {
              StringBuilder sb = new StringBuilder();
              sb.append("(");
              for (int index = 0; index < item.getArgCount(); index++) {
                  if (index > 0) {
                      sb.append(" AND ");
                  }
                  sb.append(getUpdateItemName(item.arguments().get(index)));
              }
              sb.append(")");
              return sb.toString();
          } else if (item instanceof ItemFuncNot) {
              return "(NOT " + getUpdateItemName(item.arguments().get(0)) + ")";
          } else if (item instanceof ItemBoolFunc2) {
              return getBoolFuncItemName(item);
          } else if (item.type().equals(Item.ItemType.FIELD_ITEM)) {
              String tableName = "`" + item.getTableName() + "`.`" + item.getItemName() + "`";
              if (item.getDbName() == null) {
                  return tableName;
              }
              if (item.getReferTables().size() == 0) {
                  return tableName;
              }
              PlanNode tbNode = item.getReferTables().iterator().next();
              if (!(tbNode instanceof TableNode)) {
                  return tableName;
              }
              if (!((TableNode) tbNode).getTableName().equals(item.getTableName())) {
                  return tableName;
              }
              return "`" + item.getDbName() + "`." + tableName;
          } else if (item instanceof ItemFuncIn) {
              Item a = item.arguments().get(0);
              StringBuilder sb = new StringBuilder();
              sb.append(getUpdateItemName(a));
              if (((ItemFuncIn) item).isNegate()) {
                  sb.append(" not ");
              }
              sb.append(" in (");
              for (int index = 1; index < item.arguments().size(); index++) {
                  if (index > 1) {
                      sb.append(",");
                  }
                  sb.append(getUpdateItemName(item.arguments().get(index)));
              }
              sb.append(")");
              return sb.toString();
          } else {
              return item.getItemName();
          }
      }
  
      private String getBoolFuncItemName(Item item) {
          Item a = item.arguments().get(0);
          Item b = item.arguments().get(1);
          String left = getUpdateItemName(a), right = getUpdateItemName(b);
          if (a instanceof ItemNull && !(b instanceof ItemNull)) {
              left = getUpdateItemName(b);
              right = getUpdateItemName(a);
          }
          String operator = ((ItemBoolFunc2) item).funcName();
          if (a instanceof ItemNull || b instanceof ItemNull) {
              if (item instanceof ItemFuncEqual) {
                  operator = SQLBinaryOperator.Is.getName();
              } else if (item instanceof ItemFuncNe) {
                  operator = SQLBinaryOperator.IsNot.getName();
              }
          }
          return left + " " + operator + " " + right;
      }
  
  
      public Item rebuildUpdateItem(Item item, String tableName) {
          BoolPtr reBuild = new BoolPtr(false);
          if (PlanUtil.isCmpFunc(item)) {
              Item res1 = PlanUtil.rebuildBoolSubQuery(item, 0, reBuild, new BoolPtr(false), new BoolPtr(false));
              if (res1 != null) {
                  return res1;
              }
  
              BoolPtr needExecuteNull = new BoolPtr(false);
              BoolPtr isAll = new BoolPtr(false);
              Item res2 = PlanUtil.rebuildBoolSubQuery(item, 1, reBuild, needExecuteNull, isAll);
              if (res2 != null) {
                  return res2;
              }
  
              ItemFunc func = (ItemFunc) item;
              item.setWithSubQuery(false);
              Item itemTmp = item.cloneStruct();
              for (int index = 0; index < func.getArgCount(); index++) {
                  Item arg = item.arguments().get(index);
                  if (isExplain && !(arg instanceof ItemBasicConstant) && !StringUtil.equalsIgnoreCase(tableName, arg.getTableName())) {
                      itemTmp.arguments().set(index, new ItemString(NEED_REPLACE, itemTmp.getCharsetIndex()));
                      itemTmp.setItemName(null);
                  } else if (arg instanceof ItemBasicConstant) {
                      continue;
                  } else {
                      int fieldIndex = getItemIndex(arg);
                      if (fieldIndex >= 0) {
                          itemTmp.arguments().set(index, valueItemList.get(fieldIndex));
                          itemTmp.setItemName(null);
                      }
                  }
              }
              return itemTmp;
          } else if (item instanceof ItemCondAnd || item instanceof ItemCondOr) {
              item.setWithSubQuery(false);
              Item itemTmp = item.cloneStruct();
              for (int index = 0; index < item.getArgCount(); index++) {
                  Item rebuildItem = rebuildUpdateItem(item.arguments().get(index), tableName);
                  itemTmp.arguments().set(index, rebuildItem);
                  itemTmp.setItemName(null);
              }
              return itemTmp;
          }
          return item;
      }
  
      private int getItemIndex(Item valueItem) {
          if (fieldList.size() == 1 && valueItemList.size() == 1 && valueItem.getTableName().equals(fieldList.get(0).getTableName())) {
              return 0;
          }
          for (int i = 0; i < fieldList.size(); i++) {
              if (fieldList.get(i).equals(valueItem)) {
                  return i;
              }
          }
          return -1;
      }
  
      @NotNull
      public RouteResultset buildRouteResultset() {
          this.visit();
          String sql = sqlBuilder.toString();
          LOGGER.debug("merge update——update sql:{}", sql);
          RouteResultset rrs = new RouteResultset(sql, ServerParse.UPDATE);
          rrs.setStatement(sql);
          rrs.setComplexSQL(true);
          return rrs;
      }
  }
  