/*
 * 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
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * 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.asterix.lang.sqlpp.rewrites.visitor;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.asterix.common.exceptions.CompilationException;
import org.apache.asterix.common.exceptions.ErrorCode;
import org.apache.asterix.common.functions.FunctionSignature;
import org.apache.asterix.lang.common.base.Expression;
import org.apache.asterix.lang.common.base.ILangExpression;
import org.apache.asterix.lang.common.context.Scope;
import org.apache.asterix.lang.common.expression.VariableExpr;
import org.apache.asterix.lang.common.rewrites.LangRewritingContext;
import org.apache.asterix.lang.common.struct.Identifier;
import org.apache.asterix.lang.sqlpp.clause.SelectBlock;
import org.apache.asterix.lang.sqlpp.expression.WindowExpression;
import org.apache.asterix.lang.sqlpp.util.FunctionMapUtil;
import org.apache.asterix.lang.sqlpp.util.SqlppVariableUtil;
import org.apache.asterix.lang.sqlpp.visitor.base.AbstractSqlppExpressionScopingVisitor;
import org.apache.asterix.om.functions.BuiltinFunctions;
import org.apache.hyracks.algebricks.common.utils.Pair;
import org.apache.hyracks.algebricks.core.algebra.functions.FunctionIdentifier;

/**
 * A pre-processor that adds a window field list into the AST:
 *
 * <pre>
 * FROM ... AS e, ... AS i
 * SELECT fn() OVER (...)
 * ->
 * FROM ... AS e, ... AS i
 * SELECT fn() AS w(e AS e, i AS i) OVER (...)
 * </pre>
 *
 * Also rewrites SQL-92 aggregate functions inside window expressions into SQL++ core aggregate functions
 * using the same approach as {@link SqlppGroupByAggregationSugarVisitor}
 * <br/>
 * Must be executed after {@link VariableCheckAndRewriteVisitor}
 */
public class SqlppWindowAggregationSugarVisitor extends AbstractSqlppExpressionScopingVisitor {

    private final Deque<SelectBlock> stack = new ArrayDeque<>();

    public SqlppWindowAggregationSugarVisitor(LangRewritingContext context) {
        super(context);
    }

    @Override
    public Expression visit(SelectBlock selectBlock, ILangExpression arg) throws CompilationException {
        stack.push(selectBlock);
        Expression expr = super.visit(selectBlock, arg);
        stack.pop();
        return expr;
    }

    @Override
    public Expression visit(WindowExpression winExpr, ILangExpression arg) throws CompilationException {
        if (!winExpr.hasWindowFieldList()) {
            SelectBlock selectBlock = stack.peek();
            List<Pair<Expression, Identifier>> winFieldList = createWindowFieldList(selectBlock);
            winExpr.setWindowFieldList(winFieldList);
        }

        FunctionSignature signature = winExpr.getFunctionSignature();
        FunctionIdentifier winfi = FunctionMapUtil.getInternalWindowFunction(signature);
        if (winfi != null) {
            winExpr.setFunctionSignature(new FunctionSignature(winfi));
            if (BuiltinFunctions.builtinFunctionHasProperty(winfi,
                    BuiltinFunctions.WindowFunctionProperty.HAS_LIST_ARG)) {
                if (winExpr.hasAggregateFilterExpr()) {
                    throw new CompilationException(ErrorCode.COMPILATION_ILLEGAL_USE_OF_FILTER_CLAUSE,
                            winExpr.getSourceLocation());
                }
                wrapAggregationArgument(winExpr, null);
            }
        } else if (FunctionMapUtil.isSql92AggregateFunction(signature)) {
            if (winExpr.getExprList().size() != 1) {
                // binary SQL-92 aggregate functions are not yet supported
                throw new CompilationException(ErrorCode.COMPILATION_INVALID_PARAMETER_NUMBER,
                        winExpr.getSourceLocation(), signature.getName(), winExpr.getExprList().size());
            }
            wrapAggregationArgument(winExpr, winExpr.getAggregateFilterExpr());
            winExpr.setAggregateFilterExpr(null);
            winExpr.setFunctionSignature(FunctionMapUtil.sql92ToCoreAggregateFunction(signature));
        }

        return super.visit(winExpr, arg);
    }

    private void wrapAggregationArgument(WindowExpression winExpr, Expression aggFilterExpr)
            throws CompilationException {
        VariableExpr winVar = winExpr.getWindowVar();

        Map<VariableExpr, Set<? extends Scope.SymbolAnnotation>> liveAnnotatedVars =
                scopeChecker.getCurrentScope().getLiveVariables();
        Set<VariableExpr> liveVars = liveAnnotatedVars.keySet();

        Map<VariableExpr, Set<? extends Scope.SymbolAnnotation>> localAnnotatedVars =
                scopeChecker.getCurrentScope().getLiveVariables(scopeChecker.getPrecedingScope());
        Set<VariableExpr> liveContextVars = Scope.findVariablesAnnotatedBy(localAnnotatedVars.keySet(),
                SqlppVariableAnnotation.CONTEXT_VARIABLE, localAnnotatedVars, winExpr.getSourceLocation());

        List<Pair<Expression, Identifier>> winFieldList = winExpr.getWindowFieldList();
        Map<VariableExpr, Identifier> winVarFieldMap =
                SqlppGroupByAggregationSugarVisitor.createGroupVarFieldMap(winFieldList);

        //binary SQL-92 aggregates are not yet supported, so we just need to rewrite the first argument
        List<Expression> exprList = winExpr.getExprList();
        Expression aggArgExpr = exprList.get(0);
        Expression newAggArgExpr = Sql92AggregateFunctionVisitor.wrapAggregationArgument(aggArgExpr, aggFilterExpr,
                winVar, winVarFieldMap, liveContextVars, null, liveVars, context);

        List<Expression> newExprList = new ArrayList<>(exprList);
        newExprList.set(0, newAggArgExpr);

        winExpr.setExprList(newExprList);
    }

    private List<Pair<Expression, Identifier>> createWindowFieldList(SelectBlock selectBlock) {
        List<Pair<Expression, Identifier>> fieldList = new ArrayList<>();
        if (selectBlock != null) {
            addToFieldList(fieldList, SqlppVariableUtil.getBindingVariables(selectBlock.getFromClause()));
            addToFieldList(fieldList, SqlppVariableUtil.getLetBindingVariables(selectBlock.getLetWhereList()));
            addToFieldList(fieldList, SqlppVariableUtil.getBindingVariables(selectBlock.getGroupbyClause()));
            addToFieldList(fieldList,
                    SqlppVariableUtil.getLetBindingVariables(selectBlock.getLetHavingListAfterGroupby()));
        }
        return fieldList;
    }

    private void addToFieldList(List<Pair<Expression, Identifier>> outFieldList, List<VariableExpr> varList) {
        for (VariableExpr varExpr : varList) {
            if (scopeChecker.lookupSymbol(varExpr.getVar().getValue()) != null) {
                SqlppVariableUtil.addToFieldVariableList(varExpr, outFieldList);
            }
        }
    }
}
