package util.sqlparse.visitor.mongo.desensitization;

import com.alibaba.druid.util.StringUtils;
import java.util.List;
import java.util.Map;
import org.bson.BsonArray;
import org.bson.BsonDocument;
import org.bson.BsonDouble;
import org.bson.BsonString;
import org.bson.BsonValue;
import util.sqlparse.visitor.mongo.BsonArrayNode;
import util.sqlparse.visitor.mongo.BsonBasicNode;
import util.sqlparse.visitor.mongo.BsonObjectNode;
import util.sqlparse.visitor.mongo.Identifier;
import util.sqlparse.visitor.mongo.MongoNode;
import util.sqlparse.visitor.mongo.MongoVisitor;
import util.sqlparse.visitor.mongo.ParseResult;

public class RowVisitControllerVisitor implements MongoVisitor<ParseResult> {
   private ParseResult result;
   private Map params;
   private final String tableRegex;
   private final String column;
   private final String value;
   private String limit;

   public ParseResult getResult() {
      return this.result;
   }

   public RowVisitControllerVisitor(Map params) {
      this.params = params;
      this.result = new ParseResult();
      this.result.isMatched = false;
      Map<String, List<String>> modifyTable = (Map)params.get("modifyTable");
      Map<String, String> tableKeyWord = (Map)params.get("tableKeyWord");
      this.limit = (String)params.get("limit");
      this.tableRegex = (String)tableKeyWord.keySet().iterator().next();
      this.column = (String)tableKeyWord.values().iterator().next();
      this.value = (String)((List)modifyTable.values().iterator().next()).get(0);
   }

   public ParseResult visitFind(MongoNode context) {
      BsonObjectNode ctx = (BsonObjectNode)context;
      BsonObjectNode filter = (BsonObjectNode)ctx.get(Identifier.Filter.code);
      this.generateValueFilterDoc(filter);
      this.visitLimit(ctx);
      this.result.newMongoNode = ctx;
      this.result.isMatched = true;
      return this.result;
   }

   private void visitLimit(BsonObjectNode ctx) {
      if (!StringUtils.isEmpty(this.limit)) {
         BsonBasicNode limit = (BsonBasicNode)ctx.get(Identifier.Limit.code);
         if (limit == null) {
            ctx.put(Identifier.Limit.code, (BsonValue)(new BsonDouble((double)Integer.valueOf(this.limit))));
         } else {
            double limitNum = (Double)limit.value();
            if (limitNum > (double)Integer.valueOf(this.limit) || limitNum == (double)0.0F) {
               ctx.put(Identifier.Limit.code, (BsonValue)(new BsonDouble((double)Integer.valueOf(this.limit))));
            }
         }

      }
   }

   private void generateValueFilterDoc(BsonObjectNode condition) {
      BsonDocument value = new BsonDocument();
      value.put(Identifier.NE.code, new BsonString(this.value));
      BsonDocument columnDoc = new BsonDocument();
      columnDoc.put(this.column, value);
      BsonArray andDoc = null;
      BsonArrayNode dollarAnd = (BsonArrayNode)condition.get(Identifier.And.code);
      if (dollarAnd == null) {
         andDoc = new BsonArray();
      } else {
         andDoc = (BsonArray)dollarAnd.getValue();
      }

      andDoc.add(columnDoc);
      condition.put(Identifier.And.code, (BsonValue)andDoc);
   }

   public ParseResult visitInsert(MongoNode ctx) {
      BsonArrayNode array = (BsonArrayNode)ctx;
      BsonObjectNode insert = (BsonObjectNode)array.get(array.size() - 1);
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitUpdate(MongoNode ctx) {
      BsonArrayNode array = (BsonArrayNode)ctx;
      BsonObjectNode updateFilter = (BsonObjectNode)array.get(0);
      BsonObjectNode q = (BsonObjectNode)updateFilter.get(Identifier.Q.code);
      this.generateValueFilterDoc(q);
      this.visitLimit(updateFilter);
      this.result.newMongoNode = ctx;
      this.result.isMatched = true;
      return this.result;
   }

   public ParseResult visitFindAndModify(MongoNode ctx) {
      BsonObjectNode objCtx = (BsonObjectNode)ctx;
      BsonObjectNode query = (BsonObjectNode)objCtx.get(Identifier.Query.code);
      this.generateValueFilterDoc(query);
      this.result.newMongoNode = ctx;
      this.result.isMatched = true;
      return this.result;
   }

   public ParseResult visitDelete(MongoNode ctx) {
      BsonArrayNode array = (BsonArrayNode)ctx;
      BsonObjectNode deleteFilter = (BsonObjectNode)array.get(0);
      BsonObjectNode q = (BsonObjectNode)deleteFilter.get(Identifier.Q.code);
      this.generateValueFilterDoc(q);
      this.result.newMongoNode = ctx;
      this.result.isMatched = true;
      return this.result;
   }

   public ParseResult visitCount(MongoNode ctx) {
      BsonObjectNode objCtx = (BsonObjectNode)ctx;
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitAggregate(MongoNode ctx) {
      BsonObjectNode objNode = (BsonObjectNode)ctx;
      BsonArrayNode pipline = (BsonArrayNode)objNode.get(Identifier.Pipeline.code);
      BsonArray piplineArray = (BsonArray)pipline.value();
      BsonDocument matchParent = null;

      for(BsonValue value : piplineArray) {
         for(String s : ((BsonDocument)value).keySet()) {
            if (Identifier.Match.code.equalsIgnoreCase(s)) {
               matchParent = (BsonDocument)value;
               break;
            }
         }
      }

      if (matchParent == null) {
         matchParent = new BsonDocument();
         piplineArray.add(matchParent);
      }

      BsonDocument match = (BsonDocument)matchParent.get(Identifier.Match.code);
      if (match == null) {
         match = new BsonDocument();
      }

      BsonObjectNode matchNode = new BsonObjectNode(null, match);
      this.generateValueFilterDoc(matchNode);
      matchParent.put(Identifier.Match.code, match);
      objNode.put(Identifier.Pipeline.code, (BsonValue)piplineArray);
      this.result.newMongoNode = objNode;
      this.result.isMatched = true;
      return this.result;
   }

   public ParseResult visitDistinct(MongoNode ctx) {
      BsonObjectNode objNode = (BsonObjectNode)ctx;
      BsonObjectNode query = (BsonObjectNode)objNode.get(Identifier.Query.code);
      this.generateValueFilterDoc(query);
      this.result.newMongoNode = ctx;
      this.result.isMatched = true;
      return this.result;
   }

   public ParseResult visitDrop(MongoNode ctx) {
      BsonObjectNode objNode = (BsonObjectNode)ctx;
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitCreate(MongoNode ctx) {
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitCreateIndexes(MongoNode ctx) {
      BsonObjectNode createIndexesCtx = (BsonObjectNode)ctx;
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitDropIndexes(MongoNode ctx) {
      BsonObjectNode dropIndexesCtx = (BsonObjectNode)ctx;
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitReIndex(MongoNode ctx) {
      BsonObjectNode dropIndexesCtx = (BsonObjectNode)ctx;
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitRenameCollection(MongoNode ctx) {
      BsonObjectNode renameCollectionCtx = (BsonObjectNode)ctx;
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitMapreduce(MongoNode ctx) {
      BsonObjectNode mapreduceCtx = (BsonObjectNode)ctx;
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitCreateUser(MongoNode ctx) {
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitSaslSupportedMechs(MongoNode ctx) {
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitUpdateUser(MongoNode ctx) {
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitUsersInfo(MongoNode ctx) {
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitGrantRolesToUser(MongoNode ctx) {
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitRevokeRolesFromUser(MongoNode ctx) {
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitDropUser(MongoNode ctx) {
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitCreateCollection(MongoNode ctx) {
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitCommand(MongoNode ctx) {
      this.result.newMongoNode = ctx;
      return this.result;
   }
}
