content: /**
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   * <p>
   * http://www.apache.org/licenses/LICENSE-2.0
   * <p>
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  package org.apache.atlas.query;
  
  import org.apache.atlas.query.antlr4.AtlasDSLParser.AliasExprContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.AtomEContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.CompEContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.ComparisonClauseContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.CountClauseContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.ExprContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.ExprRightContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.FromExpressionContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.FromSrcContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.GroupByExpressionContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.HasClauseContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.HasTermClauseContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.IdentifierContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.IsClauseContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.LimitOffsetContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.MaxClauseContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.MinClauseContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.OrderByExprContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.SelectExprContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.SelectExpressionContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.SingleQrySrcContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.SumClauseContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.ValueArrayContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParser.WhereClauseContext;
  import org.apache.atlas.query.antlr4.AtlasDSLParserBaseVisitor;
  import org.apache.commons.collections.CollectionUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.List;
  import java.util.Objects;
  
  public class DSLVisitor extends AtlasDSLParserBaseVisitor<Void> {
      private static final Logger LOG = LoggerFactory.getLogger(DSLVisitor.class);
  
      private static final String AND = "AND";
      private static final String OR  = "OR";
  
      private final GremlinQueryComposer gremlinQueryComposer;
  
      public DSLVisitor(GremlinQueryComposer gremlinQueryComposer) {
          this.gremlinQueryComposer = gremlinQueryComposer;
      }
  
      @Override
      public Void visitLimitOffset(LimitOffsetContext ctx) {
          LOG.debug("=> DSLVisitor.visitLimitOffset({})", ctx);
  
          gremlinQueryComposer.addLimit(ctx.limitClause().NUMBER().getText(), (ctx.offsetClause() == null ? "0" : ctx.offsetClause().NUMBER().getText()));
  
          return super.visitLimitOffset(ctx);
      }
  
      @Override
      public Void visitSelectExpr(SelectExprContext ctx) {
          LOG.debug("=> DSLVisitor.visitSelectExpr({})", ctx);
  
          // Select can have only attributes, aliased attributes or aggregate functions
  
          // Groupby attr also represent select expr, no processing is needed in that case
          // visit groupBy would handle the select expr appropriately
          if (!(ctx.getParent() instanceof GroupByExpressionContext)) {
              String[] items    = new String[ctx.selectExpression().size()];
              String[] labels   = new String[ctx.selectExpression().size()];
              int      countIdx = -1;
              int      sumIdx   = -1;
              int      minIdx   = -1;
              int      maxIdx   = -1;
  
              for (int i = 0; i < ctx.selectExpression().size(); i++) {
                  SelectExpressionContext selectExpression = ctx.selectExpression(i);
                  CountClauseContext      countClause      = selectExpression.expr().compE().countClause();
                  SumClauseContext        sumClause        = selectExpression.expr().compE().sumClause();
                  MinClauseContext        minClause        = selectExpression.expr().compE().minClause();
                  MaxClauseContext        maxClause        = selectExpression.expr().compE().maxClause();
                  IdentifierContext       identifier       = selectExpression.identifier();
  
                  labels[i] = identifier != null ? identifier.getText() : selectExpression.getText();
  
                  if (Objects.nonNull(countClause)) {
                      items[i] = "count";
                      countIdx = i;
                  } else if (Objects.nonNull(sumClause)) {
                      items[i] = sumClause.expr().getText();
                      sumIdx   = i;
                  } else if (Objects.nonNull(minClause)) {
                      items[i] = minClause.expr().getText();
                      minIdx   = i;
                  } else if (Objects.nonNull(maxClause)) {
                      items[i] = maxClause.expr().getText();
                      maxIdx   = i;
                  } else {
                      items[i] = selectExpression.expr().getText();
                  }
              }
  
              SelectClauseComposer selectClauseComposer = new SelectClauseComposer(labels, items, items, countIdx, sumIdx, minIdx, maxIdx);
  
              gremlinQueryComposer.addSelect(selectClauseComposer);
          }
  
          return super.visitSelectExpr(ctx);
      }
  
      @Override
      public Void visitOrderByExpr(OrderByExprContext ctx) {
          LOG.debug("=> DSLVisitor.visitOrderByExpr({})", ctx);
  
          // Extract the attribute from parentheses
          String text = ctx.expr().getText().replace("(", "").replace(")", "");
  
          gremlinQueryComposer.addOrderBy(text, (ctx.sortOrder() != null && ctx.sortOrder().getText().equalsIgnoreCase("desc")));
  
          return super.visitOrderByExpr(ctx);
      }
  
      @Override
      public Void visitWhereClause(WhereClauseContext ctx) {
          LOG.debug("=> DSLVisitor.visitWhereClause({})", ctx);
  
          ExprContext expr = ctx.expr();
  
          processExpr(expr, gremlinQueryComposer);
  
          return super.visitWhereClause(ctx);
      }
  
      @Override
      public Void visitFromExpression(final FromExpressionContext ctx) {
          LOG.debug("=> DSLVisitor.visitFromExpression({})", ctx);
  
          FromSrcContext   fromSrc   = ctx.fromSrc();
          AliasExprContext aliasExpr = fromSrc.aliasExpr();
  
          if (aliasExpr != null) {
              gremlinQueryComposer.addFromAlias(aliasExpr.identifier(0).getText(), aliasExpr.identifier(1).getText());
          } else {
              if (fromSrc.identifier() != null) {
                  gremlinQueryComposer.addFrom(fromSrc.identifier().getText());
              } else {
                  gremlinQueryComposer.addFrom(fromSrc.literal().getText());
              }
          }
  
          return super.visitFromExpression(ctx);
      }
  
      @Override
      public Void visitSingleQrySrc(SingleQrySrcContext ctx) {
          if (ctx.fromExpression() == null) {
              if (ctx.expr() != null && !gremlinQueryComposer.hasFromClause()) {
                  inferFromClause(ctx);
              }
  
              if (ctx.expr() != null && gremlinQueryComposer.hasFromClause()) {
                  processExpr(ctx.expr(), gremlinQueryComposer);
              }
          }
  
          return super.visitSingleQrySrc(ctx);
      }
  
      @Override
      public Void visitGroupByExpression(GroupByExpressionContext ctx) {
          LOG.debug("=> DSLVisitor.visitGroupByExpression({})", ctx);
  
          String s = ctx.selectExpr().getText();
  
          gremlinQueryComposer.addGroupBy(s);
  
          return super.visitGroupByExpression(ctx);
      }
  
      private Void visitIsClause(GremlinQueryComposer gqc, IsClauseContext ctx) {
          LOG.debug("=> DSLVisitor.visitIsClause({})", ctx);
  
          gqc.addIsA(ctx.arithE().getText(), ctx.identifier().getText());
  
          return super.visitIsClause(ctx);
      }
  
      private void visitHasClause(GremlinQueryComposer gqc, HasClauseContext ctx) {
          gqc.addFromProperty(ctx.arithE().getText(), ctx.identifier().getText());
  
          super.visitHasClause(ctx);
      }
  
      private void visitHasTermClause(GremlinQueryComposer gqc, HasTermClauseContext ctx) {
          gqc.addHasTerm(ctx.arithE().getText(), ctx.identifier().getText());
  
          super.visitHasTermClause(ctx);
      }
  
      private void inferFromClause(SingleQrySrcContext ctx) {
          if (ctx.fromExpression() != null) {
              return;
          }
  
          if (ctx.expr() != null && gremlinQueryComposer.hasFromClause()) {
              return;
          }
  
          if (ctx.expr().compE() != null && ctx.expr().compE().isClause() != null && ctx.expr().compE().isClause().arithE() != null) {
              gremlinQueryComposer.addFrom(ctx.expr().compE().isClause().arithE().getText());
  
              return;
          }
  
          if (ctx.expr().compE() != null && ctx.expr().compE().hasClause() != null && ctx.expr().compE().hasClause().arithE() != null) {
              gremlinQueryComposer.addFrom(ctx.expr().compE().hasClause().arithE().getText());
          }
  
          if (ctx.expr().compE() != null && ctx.expr().compE().hasTermClause() != null && ctx.expr().compE().hasTermClause().arithE() != null) {
              gremlinQueryComposer.addFrom(ctx.expr().compE().hasTermClause().arithE().getText());
  
              return;
          }
      }
  
      private void processExpr(final ExprContext expr, GremlinQueryComposer gremlinQueryComposer) {
          if (CollectionUtils.isNotEmpty(expr.exprRight())) {
              processExprRight(expr, gremlinQueryComposer);
          } else {
              GremlinQueryComposer original = gremlinQueryComposer.newInstance();
  
              original.addAll(gremlinQueryComposer.getQueryClauses());
  
              processExpr(expr.compE(), gremlinQueryComposer);
  
              if (gremlinQueryComposer.hasAnyTraitAttributeClause()) {
                  gremlinQueryComposer.addAll(original.getQueryClauses());
                  processExprForTrait(expr, gremlinQueryComposer);
              }
          }
      }
  
      private void processExprForTrait(final ExprContext expr, GremlinQueryComposer gremlinQueryComposer) {
          //add AND clause
          GremlinQueryComposer nestedProcessor = gremlinQueryComposer.createNestedProcessor();
  
          processExpr(expr.compE(), nestedProcessor);
  
          GremlinClauseList clauseList = nestedProcessor.getQueryClauses();
  
          if (clauseList.size() > 1) {
              gremlinQueryComposer.addAndClauses(Collections.singletonList(nestedProcessor));
          }
      }
  
      private void processExprRight(final ExprContext expr, GremlinQueryComposer gremlinQueryComposer) {
          GremlinQueryComposer       nestedProcessor = gremlinQueryComposer.createNestedProcessor();
          List<GremlinQueryComposer> nestedQueries   = new ArrayList<>();
          String                     prev            = null;
  
          // Process first expression then proceed with the others
          // expr -> compE exprRight*
          processExpr(expr.compE(), nestedProcessor);
          nestedQueries.add(nestedProcessor);
  
          // Record all processed attributes
          gremlinQueryComposer.addProcessedAttributes(nestedProcessor.getAttributesProcessed());
  
          for (ExprRightContext exprRight : expr.exprRight()) {
              nestedProcessor = gremlinQueryComposer.createNestedProcessor();
  
              // AND expression
              if (exprRight.K_AND() != null) {
                  if (OR.equalsIgnoreCase(prev)) {
                      // Change of context
                      GremlinQueryComposer orClause = nestedProcessor.createNestedProcessor();
  
                      orClause.addOrClauses(nestedQueries);
                      nestedQueries.clear();
                      nestedQueries.add(orClause);
  
                      // Record all processed attributes
                      gremlinQueryComposer.addProcessedAttributes(orClause.getAttributesProcessed());
                  }
  
                  prev = AND;
              }
  
              // OR expression
              if (exprRight.K_OR() != null) {
                  if (AND.equalsIgnoreCase(prev)) {
                      // Change of context
                      GremlinQueryComposer andClause = nestedProcessor.createNestedProcessor();
  
                      andClause.addAndClauses(nestedQueries);
                      nestedQueries.clear();
                      nestedQueries.add(andClause);
  
                      // Record all processed attributes
                      gremlinQueryComposer.addProcessedAttributes(andClause.getAttributesProcessed());
                  }
  
                  prev = OR;
              }
  
              processExpr(exprRight.compE(), nestedProcessor);
              nestedQueries.add(nestedProcessor);
  
              // Record all processed attributes
              gremlinQueryComposer.addProcessedAttributes(nestedProcessor.getAttributesProcessed());
          }
  
          if (AND.equalsIgnoreCase(prev)) {
              gremlinQueryComposer.addAndClauses(nestedQueries);
          } else if (OR.equalsIgnoreCase(prev)) {
              gremlinQueryComposer.addOrClauses(nestedQueries);
          }
      }
  
      private void processExpr(final CompEContext compE, final GremlinQueryComposer gremlinQueryComposer) {
          if (compE != null) {
              IsClauseContext      isClause      = compE.isClause();
              HasClauseContext     hasClause     = compE.hasClause();
              HasTermClauseContext hasTermClause = compE.hasTermClause();
  
              if (isClause != null) {
                  visitIsClause(gremlinQueryComposer, isClause);
              }
  
              if (hasClause != null) {
                  visitHasClause(gremlinQueryComposer, hasClause);
              }
  
              if (hasTermClause != null) {
                  visitHasTermClause(gremlinQueryComposer, hasTermClause);
              }
  
              if (isClause == null && hasClause == null && hasTermClause == null) {
                  ComparisonClauseContext comparisonClause = compE.comparisonClause();
  
                  // The nested expression might have ANDs/ORs
                  if (comparisonClause == null) {
                      ExprContext exprContext = compE.arithE().multiE().atomE().expr();
  
                      // Only extract comparison clause if there are no nested exprRight clauses
                      if (CollectionUtils.isEmpty(exprContext.exprRight())) {
                          comparisonClause = exprContext.compE().comparisonClause();
                      }
                  }
  
                  if (comparisonClause != null) {
                      String       lhs      = comparisonClause.arithE(0).getText();
                      AtomEContext atomECtx = comparisonClause.arithE(1).multiE().atomE();
                      String       op;
                      String       rhs;
  
                      if (atomECtx.literal() == null || (atomECtx.literal() != null && atomECtx.literal().valueArray() == null)) {
                          op  = comparisonClause.operator().getText().toUpperCase();
                          rhs = comparisonClause.arithE(1).getText();
                      } else {
                          op  = "in";
                          rhs = getInClause(atomECtx);
                      }
  
                      gremlinQueryComposer.addWhere(lhs, op, rhs);
                  } else {
                      processExpr(compE.arithE().multiE().atomE().expr(), gremlinQueryComposer);
                  }
              }
          }
      }
  
      private String getInClause(AtomEContext atomEContext) {
          StringBuilder     sb                = new StringBuilder();
          ValueArrayContext valueArrayContext = atomEContext.literal().valueArray();
          int               startIdx          = 1;
          int               endIdx            = valueArrayContext.children.size() - 1;
  
          for (int i = startIdx; i < endIdx; i++) {
              sb.append(valueArrayContext.getChild(i));
          }
  
          return sb.toString();
      }
  }
  