/*
 * 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.optimizer.rules;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

import org.apache.asterix.dataflow.data.common.TypeResolverUtil;
import org.apache.asterix.lang.common.util.FunctionUtil;
import org.apache.asterix.om.functions.BuiltinFunctions;
import org.apache.asterix.om.typecomputer.base.TypeCastUtils;
import org.apache.asterix.om.types.IAType;
import org.apache.commons.lang3.mutable.Mutable;
import org.apache.commons.lang3.mutable.MutableObject;
import org.apache.hyracks.algebricks.common.exceptions.AlgebricksException;
import org.apache.hyracks.algebricks.common.utils.Triple;
import org.apache.hyracks.algebricks.core.algebra.base.ILogicalExpression;
import org.apache.hyracks.algebricks.core.algebra.base.ILogicalOperator;
import org.apache.hyracks.algebricks.core.algebra.base.IOptimizationContext;
import org.apache.hyracks.algebricks.core.algebra.base.LogicalOperatorTag;
import org.apache.hyracks.algebricks.core.algebra.base.LogicalVariable;
import org.apache.hyracks.algebricks.core.algebra.expressions.IVariableTypeEnvironment;
import org.apache.hyracks.algebricks.core.algebra.expressions.ScalarFunctionCallExpression;
import org.apache.hyracks.algebricks.core.algebra.expressions.VariableReferenceExpression;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.AssignOperator;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.UnionAllOperator;
import org.apache.hyracks.algebricks.core.rewriter.base.IAlgebraicRewriteRule;

/**
 * This rule injects type casts for inputs of a UnionAllOperator if those
 * inputs have heterogeneous types.
 */
public class InjectTypeCastForUnionRule implements IAlgebraicRewriteRule {

    @Override
    public boolean rewritePost(Mutable<ILogicalOperator> opRef, IOptimizationContext context)
            throws AlgebricksException {
        ILogicalOperator op = opRef.getValue();
        if (op.getOperatorTag() != LogicalOperatorTag.UNIONALL) {
            return false;
        }
        UnionAllOperator unionAllOperator = (UnionAllOperator) op;
        // Injects casts to the first and second input branch of the UnionAllOperator.
        return injectCast(unionAllOperator, 0, context) || injectCast(unionAllOperator, 1, context);
    }

    // Injects a type cast function on one input (indicated by childIndex) of the union all operator if necessary.
    private boolean injectCast(UnionAllOperator op, int childIndex, IOptimizationContext context)
            throws AlgebricksException {
        // Gets the type environments for the union all operator and its child operator with the right child index.
        IVariableTypeEnvironment env = context.getOutputTypeEnvironment(op);
        Mutable<ILogicalOperator> branchOpRef = op.getInputs().get(childIndex);
        IVariableTypeEnvironment childEnv = context.getOutputTypeEnvironment(branchOpRef.getValue());

        // The two lists are used for the assign operator that calls cast functions.
        List<LogicalVariable> varsToCast = new ArrayList<>();
        List<Mutable<ILogicalExpression>> castFunctionsForLeft = new ArrayList<>();

        // Iterate through all triples.
        List<Triple<LogicalVariable, LogicalVariable, LogicalVariable>> triples = op.getVariableMappings();
        for (Triple<LogicalVariable, LogicalVariable, LogicalVariable> triple : triples) {
            LogicalVariable producedVar = triple.third;
            IAType producedType = (IAType) env.getVarType(producedVar);
            LogicalVariable varToCast = childIndex == 0 ? triple.first : triple.second;
            IAType inputType = (IAType) childEnv.getVarType(varToCast);
            if (!TypeResolverUtil.needsCast(producedType, inputType)) {
                // Continues to the next triple if no cast is neeeded.
                continue;
            }
            LogicalVariable castedVar = context.newVar();
            // Resets triple variables to new variables that bind to the results of type casting.
            triple.first = childIndex == 0 ? castedVar : triple.first;
            triple.second = childIndex > 0 ? castedVar : triple.second;
            ScalarFunctionCallExpression castFunc = new ScalarFunctionCallExpression(
                    FunctionUtil.getFunctionInfo(BuiltinFunctions.CAST_TYPE), new ArrayList<>(Collections
                            .singletonList(new MutableObject<>(new VariableReferenceExpression(varToCast)))));
            TypeCastUtils.setRequiredAndInputTypes(castFunc, producedType, inputType);

            // Adds the variable and function expression into lists, for the assign operator.
            varsToCast.add(castedVar);
            castFunctionsForLeft.add(new MutableObject<>(castFunc));
        }
        if (castFunctionsForLeft.isEmpty()) {
            return false;
        }
        // Injects an assign operator to perform type casts.
        AssignOperator assignOp = new AssignOperator(varsToCast, castFunctionsForLeft);
        assignOp.getInputs().add(new MutableObject<>(branchOpRef.getValue()));
        branchOpRef.setValue(assignOp);
        context.computeAndSetTypeEnvironmentForOperator(assignOp);

        // Returns true to indicate that rewriting happens.
        return true;
    }

}
