package util.sqlparse.visitor.mongo.desensitization;

import java.util.ArrayList;
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.BsonJsScopeNode;
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 FieldReplaceVisitor implements MongoVisitor<ParseResult> {
   private ParseResult result;
   private Map params;
   private ArrayList<ParseResult.Column> list;
   private String replaceSchema;
   private String replaceTable;
   private String replaceColumn;
   private String expression;

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

   public FieldReplaceVisitor(Map params, ArrayList<ParseResult.Column> list) {
      this.params = params;
      this.result = new ParseResult();
      this.result.isMatched = false;
      this.list = list;
      this.replaceSchema = (String)params.get("replaceSchema");
      this.replaceTable = (String)params.get("replaceTable");
      this.replaceColumn = (String)params.get("replaceColumn");
      this.expression = (String)params.get("expression");
   }

   public ParseResult visitFind(MongoNode context) {
      BsonObjectNode ctx = (BsonObjectNode)context;
      MongoNode find = ctx.get(Identifier.Find.code);
      String collection = find.value().toString();
      BsonObjectNode filter = (BsonObjectNode)ctx.get(Identifier.Filter.code);
      BsonObjectNode projection = (BsonObjectNode)ctx.get(Identifier.Projection.code);
      BsonObjectNode aggregateCtx = new BsonObjectNode();
      BsonDocument aggregateDoc = new BsonDocument();
      aggregateDoc.put(Identifier.Aggregate.code, new BsonString(collection));
      BsonArray pipline = new BsonArray();
      this.visitPorject(projection, pipline);
      this.visitPipline(ctx, pipline);
      if (filter != null) {
         BsonDocument temp = new BsonDocument();
         temp.put(Identifier.Match.code, filter.getValue());
         pipline.add(temp);
      }

      aggregateDoc.put(Identifier.Pipeline.code, pipline);
      this.mergeDocument(aggregateDoc, ctx.getValue(), Identifier.Explain, Identifier.AllowDiskUse, Identifier.Cursor, Identifier.MaxTimeMS, Identifier.BypassDocumentValidation, Identifier.ReadConcern, Identifier.Collation, Identifier.Hint, Identifier.Comment, Identifier.WriteConcern, Identifier.Let);
      BsonValue explain = aggregateDoc.get(Identifier.Explain);
      BsonValue cursor = aggregateDoc.get(Identifier.Cursor);
      if (explain == null && cursor == null) {
         BsonDocument cursorDoc = new BsonDocument();
         aggregateDoc.put(Identifier.Cursor.code, cursorDoc);
      }

      this.mergeDocument(aggregateDoc, ctx.getValue(), Identifier.LSID, Identifier.DB);
      aggregateCtx.setValue(aggregateDoc);
      this.result.newMongoNode = aggregateCtx;
      return this.result;
   }

   private void visitPorject(BsonObjectNode projection, BsonArray array) {
      BsonDocument project = null;
      if (projection != null) {
         project = projection.getValue().clone();
      }

      this.replaceExpression(array, project);
   }

   private void visitPipline(BsonObjectNode ctx, BsonArray array) {
      this.visitItem(ctx, array, Identifier.Sort, Identifier.SORT);
      this.visitItem(ctx, array, Identifier.Limit, Identifier.DollarLimit);
      this.visitItem(ctx, array, Identifier.Skip, Identifier.DollarSkip);
   }

   private void visitItem(BsonObjectNode ctx, BsonArray array, Identifier source, Identifier target) {
      MongoNode item = ctx.get(source.code);
      if (item != null) {
         BsonDocument sortDoc = new BsonDocument();
         if (item instanceof BsonObjectNode) {
            sortDoc.put(target.code, ((BsonObjectNode)item).getValue());
         } else if (item instanceof BsonBasicNode) {
            sortDoc.put(target.code, ((BsonBasicNode)item).getValue());
         } else if (item instanceof BsonArrayNode) {
            sortDoc.put(target.code, ((BsonArrayNode)item).getValue());
         } else if (item instanceof BsonJsScopeNode) {
            sortDoc.put(target.code, ((BsonJsScopeNode)item).getValue());
         } else {
            sortDoc.put(target.code, ((BsonObjectNode)item).getValue());
         }

         array.add(sortDoc);
      }

   }

   private void replaceExpression(BsonArray array, BsonDocument oldPreject) {
      for(ParseResult.Column column : this.list) {
         String newExpression = this.expression.replaceAll("\\$\\{value}", "\\$" + column.name + "");

         for(BsonValue value : BsonArray.parse(newExpression)) {
            if (value instanceof BsonDocument) {
               BsonDocument doc = (BsonDocument)value;
               BsonDocument project = (BsonDocument)doc.get("$project");
               if (project != null) {
                  String key = "$" + column.name;
                  BsonValue value2 = project.get(key);
                  if (value2 != null) {
                     project.remove(key);
                     if (oldPreject != null) {
                        for(String s : oldPreject.keySet()) {
                           BsonValue value1 = oldPreject.get(s);
                           if (!(value1 instanceof BsonDocument) && !(value1 instanceof BsonArray)) {
                              project.put(s, value1);
                           } else {
                              project.put(s, new BsonDouble((double)1.0F));
                           }
                        }
                     }

                     project.put(column.name, value2);
                  }
               }

               array.add(doc);
            }
         }
      }

   }

   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 updateCtx = (BsonObjectNode)array.get(array.size() - 1);
      this.result.newMongoNode = ctx;
      return this.result;
   }

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

   public ParseResult visitDelete(MongoNode ctx) {
      BsonArrayNode array = (BsonArrayNode)ctx;
      BsonObjectNode collectionCtx = (BsonObjectNode)array.get(array.size() - 1);
      this.result.newMongoNode = ctx;
      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 project = null;

      for(int i = piplineArray.size() - 1; i >= 0; --i) {
         BsonValue value = piplineArray.get(i);
         if (value instanceof BsonDocument) {
            BsonDocument doc = (BsonDocument)value;
            project = (BsonDocument)doc.get("$project");
            if (project != null) {
               break;
            }
         }
      }

      this.replaceExpression(piplineArray, project);
      objNode.put(Identifier.Pipeline.code, (BsonValue)piplineArray);
      this.result.newMongoNode = ctx;
      return this.result;
   }

   public ParseResult visitDistinct(MongoNode ctx) {
      BsonObjectNode objNode = (BsonObjectNode)ctx;
      this.result.newMongoNode = ctx;
      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;
   }

   private void mergeDocument(BsonDocument target, BsonDocument source, Identifier... ids) {
      for(Identifier id : ids) {
         BsonValue value = source.get(id.code);
         if (value != null) {
            target.put(id.code, value);
         }
      }

   }
}
