/*
 * 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.HashSet;
import java.util.LinkedList;
import java.util.List;

import org.apache.asterix.algebra.base.OperatorAnnotation;
import org.apache.asterix.common.config.DatasetConfig.DatasetType;
import org.apache.asterix.common.exceptions.CompilationException;
import org.apache.asterix.common.exceptions.ErrorCode;
import org.apache.asterix.metadata.declared.DataSource;
import org.apache.asterix.metadata.declared.DataSourceId;
import org.apache.asterix.metadata.declared.MetadataProvider;
import org.apache.asterix.metadata.entities.Dataset;
import org.apache.asterix.metadata.entities.Index;
import org.apache.asterix.metadata.utils.DatasetUtil;
import org.apache.asterix.om.base.AInt32;
import org.apache.asterix.om.base.AString;
import org.apache.asterix.om.base.IAObject;
import org.apache.asterix.om.constants.AsterixConstantValue;
import org.apache.asterix.om.types.ARecordType;
import org.apache.asterix.om.types.ATypeTag;
import org.apache.asterix.om.types.IAType;
import org.apache.asterix.om.utils.ConstantExpressionUtil;
import org.apache.asterix.optimizer.base.AnalysisUtil;
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.Pair;
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.ILogicalPlan;
import org.apache.hyracks.algebricks.core.algebra.base.IOptimizationContext;
import org.apache.hyracks.algebricks.core.algebra.base.LogicalExpressionTag;
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.AbstractFunctionCallExpression;
import org.apache.hyracks.algebricks.core.algebra.expressions.ConstantExpression;
import org.apache.hyracks.algebricks.core.algebra.expressions.VariableReferenceExpression;
import org.apache.hyracks.algebricks.core.algebra.metadata.IDataSource;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.AbstractLogicalOperator;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.AbstractOperatorWithNestedPlans;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.AssignOperator;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.DataSourceScanOperator;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.GroupByOperator;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.NestedTupleSourceOperator;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.visitors.VariableUtilities;
import org.apache.hyracks.algebricks.core.algebra.util.OperatorPropertiesUtil;
import org.apache.hyracks.algebricks.core.rewriter.base.IAlgebraicRewriteRule;

public class PushFieldAccessRule implements IAlgebraicRewriteRule {

    @Override
    public boolean rewritePre(Mutable<ILogicalOperator> opRef, IOptimizationContext context) {
        return false;
    }

    @Override
    public boolean rewritePost(Mutable<ILogicalOperator> opRef, IOptimizationContext context)
            throws AlgebricksException {
        AbstractLogicalOperator op = (AbstractLogicalOperator) opRef.getValue();
        if (context.checkIfInDontApplySet(this, op)) {
            return false;
        }
        if (op.getOperatorTag() != LogicalOperatorTag.ASSIGN) {
            return false;
        }
        AssignOperator access = (AssignOperator) op;
        ILogicalExpression expr = getFirstExpr(access);
        String finalAnnot;
        if (AnalysisUtil.isAccessToFieldRecord(expr)) {
            finalAnnot = OperatorAnnotation.PUSHED_FIELD_ACCESS;
        } else if (AnalysisUtil.isRunnableAccessToFieldRecord(expr)) {
            finalAnnot = OperatorAnnotation.PUSHED_RUNNABLE_FIELD_ACCESS;
        } else {
            return false;
        }
        return propagateFieldAccessRec(opRef, context, finalAnnot);
    }

    private boolean isAccessToIndexedField(AssignOperator assign, IOptimizationContext context)
            throws AlgebricksException {
        AbstractFunctionCallExpression accessFun =
                (AbstractFunctionCallExpression) assign.getExpressions().get(0).getValue();
        ILogicalExpression e0 = accessFun.getArguments().get(0).getValue();
        if (e0.getExpressionTag() != LogicalExpressionTag.VARIABLE) {
            return false;
        }
        LogicalVariable var = ((VariableReferenceExpression) e0).getVariableReference();
        if (context.findPrimaryKey(var) == null) {
            // not referring to a dataset record
            return false;
        }
        AbstractLogicalOperator op = assign;
        while (op.getInputs().size() == 1 && op.getOperatorTag() != LogicalOperatorTag.DATASOURCESCAN) {
            op = (AbstractLogicalOperator) op.getInputs().get(0).getValue();
        }
        if (op.getOperatorTag() != LogicalOperatorTag.DATASOURCESCAN) {
            return false;
        }
        DataSourceScanOperator scan = (DataSourceScanOperator) op;
        LogicalVariable recVar = scan.getVariables().get(scan.getVariables().size() - 1);
        if (recVar != var) {
            return false;
        }
        MetadataProvider mp = (MetadataProvider) context.getMetadataProvider();
        DataSourceId asid = ((IDataSource<DataSourceId>) scan.getDataSource()).getId();

        Dataset dataset = mp.findDataset(asid.getDataverseName(), asid.getDatasourceName());
        if (dataset == null) {
            throw new CompilationException(ErrorCode.UNKNOWN_DATASET_IN_DATAVERSE, scan.getSourceLocation(),
                    asid.getDatasourceName(), asid.getDataverseName());
        }
        if (dataset.getDatasetType() != DatasetType.INTERNAL) {
            return false;
        }
        final Integer pos = ConstantExpressionUtil.getIntConstant(accessFun.getArguments().get(1).getValue());
        if (pos != null) {
            String tName = dataset.getItemTypeName();
            IAType t = mp.findType(dataset.getItemTypeDataverseName(), tName);
            if (t.getTypeTag() != ATypeTag.OBJECT) {
                return false;
            }
            ARecordType rt = (ARecordType) t;
            if (pos >= rt.getFieldNames().length) {
                return false;
            }
        }

        List<Index> datasetIndexes = mp.getDatasetIndexes(dataset.getDataverseName(), dataset.getDatasetName());
        boolean hasSecondaryIndex = false;
        for (Index index : datasetIndexes) {
            if (index.isSecondaryIndex()) {
                hasSecondaryIndex = true;
                break;
            }
        }

        return hasSecondaryIndex;
    }

    private boolean tryingToPushThroughSelectionWithSameDataSource(AssignOperator access, AbstractLogicalOperator op2) {
        if (op2.getOperatorTag() != LogicalOperatorTag.SELECT) {
            return false;
        }
        ILogicalExpression e1 = (ILogicalExpression) access.getAnnotations().get(OperatorAnnotation.FIELD_ACCESS);
        if (e1 == null) {
            return false;
        }
        ILogicalExpression e2 = (ILogicalExpression) op2.getAnnotations().get(OperatorAnnotation.FIELD_ACCESS);
        if (e2 == null) {
            return false;
        }
        return e1.equals(e2);
    }

    private boolean propagateFieldAccessRec(Mutable<ILogicalOperator> opRef, IOptimizationContext context,
            String finalAnnot) throws AlgebricksException {
        AssignOperator access = (AssignOperator) opRef.getValue();
        Mutable<ILogicalOperator> opRef2 = access.getInputs().get(0);
        AbstractLogicalOperator op2 = (AbstractLogicalOperator) opRef2.getValue();
        // If it's not an indexed field, it is pushed so that scan can be
        // rewritten into index search.
        if (op2.getOperatorTag() == LogicalOperatorTag.PROJECT || context.checkAndAddToAlreadyCompared(access, op2)
                && !(op2.getOperatorTag() == LogicalOperatorTag.SELECT && isAccessToIndexedField(access, context))) {
            return false;
        }
        Object annotation = op2.getAnnotations().get(OperatorPropertiesUtil.MOVABLE);
        if (annotation != null && !((Boolean) annotation)) {
            return false;
        }
        if (tryingToPushThroughSelectionWithSameDataSource(access, op2)) {
            return false;
        }
        if (testAndModifyRedundantOp(access, op2)) {
            propagateFieldAccessRec(opRef2, context, finalAnnot);
            return true;
        }
        List<LogicalVariable> usedInAccess = new LinkedList<>();
        VariableUtilities.getUsedVariables(access, usedInAccess);

        List<LogicalVariable> produced2 = new LinkedList<>();
        if (op2.getOperatorTag() == LogicalOperatorTag.GROUP) {
            VariableUtilities.getLiveVariables(op2, produced2);
        } else {
            VariableUtilities.getProducedVariables(op2, produced2);
        }
        boolean pushItDown = false;
        List<LogicalVariable> inter = new ArrayList<>(usedInAccess);
        if (inter.isEmpty()) { // ground value
            return false;
        }
        inter.retainAll(produced2);
        if (inter.isEmpty()) {
            pushItDown = true;
        } else if (op2.getOperatorTag() == LogicalOperatorTag.GROUP) {
            GroupByOperator g = (GroupByOperator) op2;
            List<Pair<LogicalVariable, LogicalVariable>> varMappings = new ArrayList<>();
            for (Pair<LogicalVariable, Mutable<ILogicalExpression>> p : g.getDecorList()) {
                ILogicalExpression e = p.second.getValue();
                if (e.getExpressionTag() == LogicalExpressionTag.VARIABLE) {
                    LogicalVariable decorVar = GroupByOperator.getDecorVariable(p);
                    if (inter.contains(decorVar)) {
                        inter.remove(decorVar);
                        LogicalVariable v1 = ((VariableReferenceExpression) e).getVariableReference();
                        varMappings.add(new Pair<>(decorVar, v1));
                    }
                }
            }
            if (inter.isEmpty()) {
                boolean changed = false;
                for (Pair<LogicalVariable, LogicalVariable> m : varMappings) {
                    LogicalVariable v2 = context.newVar();
                    LogicalVariable oldVar = access.getVariables().get(0);
                    VariableReferenceExpression v2Ref = new VariableReferenceExpression(v2);
                    v2Ref.setSourceLocation(g.getSourceLocation());
                    g.getDecorList().add(new Pair<LogicalVariable, Mutable<ILogicalExpression>>(oldVar,
                            new MutableObject<ILogicalExpression>(v2Ref)));
                    changed = true;
                    access.getVariables().set(0, v2);
                    VariableUtilities.substituteVariables(access, m.first, m.second, context);
                }
                if (changed) {
                    context.computeAndSetTypeEnvironmentForOperator(g);
                }
                usedInAccess.clear();
                VariableUtilities.getUsedVariables(access, usedInAccess);
                pushItDown = true;
            }
        }
        if (pushItDown) {
            if (op2.getOperatorTag() == LogicalOperatorTag.NESTEDTUPLESOURCE) {
                Mutable<ILogicalOperator> childOfSubplan =
                        ((NestedTupleSourceOperator) op2).getDataSourceReference().getValue().getInputs().get(0);
                pushAccessDown(opRef, op2, childOfSubplan, context, finalAnnot);
                return true;
            }
            if (op2.getInputs().size() == 1 && !op2.hasNestedPlans()) {
                pushAccessDown(opRef, op2, op2.getInputs().get(0), context, finalAnnot);
                return true;
            } else {
                for (Mutable<ILogicalOperator> inp : op2.getInputs()) {
                    HashSet<LogicalVariable> v2 = new HashSet<>();
                    VariableUtilities.getLiveVariables(inp.getValue(), v2);
                    if (v2.containsAll(usedInAccess)) {
                        pushAccessDown(opRef, op2, inp, context, finalAnnot);
                        return true;
                    }
                }
            }
            if (op2.hasNestedPlans()) {
                AbstractOperatorWithNestedPlans nestedOp = (AbstractOperatorWithNestedPlans) op2;
                for (ILogicalPlan plan : nestedOp.getNestedPlans()) {
                    for (Mutable<ILogicalOperator> root : plan.getRoots()) {
                        HashSet<LogicalVariable> v2 = new HashSet<>();
                        VariableUtilities.getLiveVariables(root.getValue(), v2);
                        if (v2.containsAll(usedInAccess)) {
                            pushAccessDown(opRef, op2, root, context, finalAnnot);
                            return true;
                        }
                    }
                }
            }
            throw new CompilationException(ErrorCode.COMPILATION_ERROR, access.getSourceLocation(),
                    "Field access " + access.getExpressions().get(0).getValue()
                            + " does not correspond to any input of operator " + op2);
        } else {
            // Check if the accessed field is not one of the partitioning key
            // fields. If yes, we can equate the two variables.
            if (op2.getOperatorTag() == LogicalOperatorTag.DATASOURCESCAN) {
                DataSourceScanOperator scan = (DataSourceScanOperator) op2;
                int n = scan.getVariables().size();
                LogicalVariable scanRecordVar = scan.getVariables().get(n - 1);

                IDataSource<DataSourceId> dataSource = (IDataSource<DataSourceId>) scan.getDataSource();
                byte dsType = ((DataSource) dataSource).getDatasourceType();
                if (dsType != DataSource.Type.INTERNAL_DATASET && dsType != DataSource.Type.EXTERNAL_DATASET) {
                    return false;
                }
                DataSourceId asid = dataSource.getId();
                MetadataProvider mp = (MetadataProvider) context.getMetadataProvider();
                Dataset dataset = mp.findDataset(asid.getDataverseName(), asid.getDatasourceName());
                if (dataset == null) {
                    throw new CompilationException(ErrorCode.UNKNOWN_DATASET_IN_DATAVERSE, scan.getSourceLocation(),
                            asid.getDatasourceName(), asid.getDataverseName());
                }
                if (dataset.getDatasetType() != DatasetType.INTERNAL) {
                    setAsFinal(access, context, finalAnnot);
                    return false;
                }

                String tName = dataset.getItemTypeName();
                IAType t = mp.findType(dataset.getItemTypeDataverseName(), tName);
                if (t.getTypeTag() != ATypeTag.OBJECT) {
                    return false;
                }
                ARecordType rt = (ARecordType) t;
                Pair<ILogicalExpression, List<String>> fieldPathAndVar = getFieldExpression(access, rt);
                ILogicalExpression e0 = fieldPathAndVar.first;
                LogicalExpressionTag tag = e0.getExpressionTag();
                if (tag == LogicalExpressionTag.VARIABLE) {
                    VariableReferenceExpression varRef = (VariableReferenceExpression) e0;
                    if (varRef.getVariableReference() == scanRecordVar) {
                        int p = DatasetUtil.getPositionOfPartitioningKeyField(dataset, fieldPathAndVar.second);
                        if (p < 0) { // not one of the partitioning fields
                            setAsFinal(access, context, finalAnnot);
                            return false;
                        }
                        LogicalVariable keyVar = scan.getVariables().get(p);
                        VariableReferenceExpression keyVarRef = new VariableReferenceExpression(keyVar);
                        keyVarRef.setSourceLocation(varRef.getSourceLocation());
                        access.getExpressions().get(0).setValue(keyVarRef);
                        return true;

                    }
                }
            }
            setAsFinal(access, context, finalAnnot);
            return false;
        }
    }

    private Pair<ILogicalExpression, List<String>> getFieldExpression(AssignOperator access, ARecordType rt)
            throws AlgebricksException {
        LinkedList<String> fieldPath = new LinkedList<>();
        ILogicalExpression e0 = access.getExpressions().get(0).getValue();
        while (AnalysisUtil.isAccessToFieldRecord(e0)) {
            ILogicalExpression e1 = ((AbstractFunctionCallExpression) e0).getArguments().get(1).getValue();
            if (e1.getExpressionTag() != LogicalExpressionTag.CONSTANT) {
                return new Pair<>(null, null);
            }
            ConstantExpression ce = (ConstantExpression) e1;
            IAObject obj = ((AsterixConstantValue) ce.getValue()).getObject();
            String fldName;
            if (obj.getType().getTypeTag() == ATypeTag.STRING) {
                fldName = ((AString) obj).getStringValue();
            } else {
                int pos = ((AInt32) obj).getIntegerValue();
                if (pos >= rt.getFieldNames().length) {
                    return new Pair<>(null, null);
                }
                fldName = rt.getFieldNames()[pos];
            }
            fieldPath.addFirst(fldName);
            e0 = ((AbstractFunctionCallExpression) e0).getArguments().get(0).getValue();

        }
        return new Pair<>(e0, fieldPath);
    }

    private void setAsFinal(ILogicalOperator access, IOptimizationContext context, String finalAnnot) {
        access.getAnnotations().put(finalAnnot, true);
        context.addToDontApplySet(this, access);
    }

    private boolean testAndModifyRedundantOp(AssignOperator access, AbstractLogicalOperator op2) {
        if (op2.getOperatorTag() != LogicalOperatorTag.ASSIGN) {
            return false;
        }
        AssignOperator a2 = (AssignOperator) op2;
        ILogicalExpression accessExpr0 = getFirstExpr(access);
        if (accessExpr0.equals(getFirstExpr(a2))) {
            VariableReferenceExpression varRef = new VariableReferenceExpression(a2.getVariables().get(0));
            varRef.setSourceLocation(accessExpr0.getSourceLocation());
            access.getExpressions().get(0).setValue(varRef);
            return true;
        } else {
            return false;
        }
    }

    // indirect recursivity with propagateFieldAccessRec
    private void pushAccessDown(Mutable<ILogicalOperator> fldAccessOpRef, ILogicalOperator op2,
            Mutable<ILogicalOperator> inputOfOp2, IOptimizationContext context, String finalAnnot)
            throws AlgebricksException {
        ILogicalOperator fieldAccessOp = fldAccessOpRef.getValue();
        fldAccessOpRef.setValue(op2);
        List<Mutable<ILogicalOperator>> faInpList = fieldAccessOp.getInputs();
        faInpList.clear();
        faInpList.add(new MutableObject<>(inputOfOp2.getValue()));
        inputOfOp2.setValue(fieldAccessOp);
        // typing
        context.computeAndSetTypeEnvironmentForOperator(fieldAccessOp);
        context.computeAndSetTypeEnvironmentForOperator(op2);
        propagateFieldAccessRec(inputOfOp2, context, finalAnnot);
    }

    private ILogicalExpression getFirstExpr(AssignOperator assign) {
        return assign.getExpressions().get(0).getValue();
    }

}
