/*
 * Copyright 2018 Confluent Inc.
 *
 * Licensed under the Confluent Community License (the "License"); you may not use
 * this file except in compliance with the License.  You may obtain a copy of the
 * License at
 *
 * http://www.confluent.io/confluent-community-license
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OF ANY KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations under the License.
 */

package io.confluent.ksql.util;

import io.confluent.ksql.function.FunctionRegistry;
import io.confluent.ksql.parser.tree.Expression;
import io.confluent.ksql.parser.tree.ExpressionRewriter;
import io.confluent.ksql.parser.tree.ExpressionTreeRewriter;
import io.confluent.ksql.parser.tree.FunctionCall;
import io.confluent.ksql.parser.tree.QualifiedName;
import io.confluent.ksql.parser.tree.QualifiedNameReference;
import java.util.ArrayList;
import java.util.List;

public class AggregateExpressionRewriter extends ExpressionRewriter<Void> {

  public static final String AGGREGATE_FUNCTION_VARIABLE_PREFIX = "KSQL_AGG_VARIABLE_";
  private int aggVariableIndex = 0;
  private final FunctionRegistry functionRegistry;

  public AggregateExpressionRewriter(final FunctionRegistry functionRegistry) {
    this.functionRegistry = functionRegistry;
  }

  @Override
  public Expression rewriteFunctionCall(final FunctionCall node, final Void context,
                                        final ExpressionTreeRewriter<Void> treeRewriter) {
    final String functionName = node.getName().getSuffix();
    if (functionRegistry.isAggregate(functionName)) {
      final String aggVarName = AGGREGATE_FUNCTION_VARIABLE_PREFIX + aggVariableIndex;
      aggVariableIndex++;
      return new QualifiedNameReference(QualifiedName.of(aggVarName));
    } else {
      final List<Expression> arguments = new ArrayList<>();
      for (final Expression argExpression: node.getArguments()) {
        arguments.add(treeRewriter.rewrite(argExpression, context));
      }
      return new FunctionCall(node.getName(), arguments);
    }
  }
}
