/*
 * 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.common.visitor;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import org.apache.asterix.common.exceptions.CompilationException;
import org.apache.asterix.common.exceptions.ErrorCode;
import org.apache.asterix.lang.common.base.Expression;
import org.apache.asterix.lang.common.base.Expression.Kind;
import org.apache.asterix.lang.common.base.ILangExpression;
import org.apache.asterix.lang.common.base.IVisitorExtension;
import org.apache.asterix.lang.common.clause.GroupbyClause;
import org.apache.asterix.lang.common.clause.LetClause;
import org.apache.asterix.lang.common.clause.LimitClause;
import org.apache.asterix.lang.common.clause.OrderbyClause;
import org.apache.asterix.lang.common.clause.WhereClause;
import org.apache.asterix.lang.common.expression.CallExpr;
import org.apache.asterix.lang.common.expression.FieldAccessor;
import org.apache.asterix.lang.common.expression.FieldBinding;
import org.apache.asterix.lang.common.expression.GbyVariableExpressionPair;
import org.apache.asterix.lang.common.expression.IfExpr;
import org.apache.asterix.lang.common.expression.IndexAccessor;
import org.apache.asterix.lang.common.expression.ListConstructor;
import org.apache.asterix.lang.common.expression.ListSliceExpression;
import org.apache.asterix.lang.common.expression.LiteralExpr;
import org.apache.asterix.lang.common.expression.OperatorExpr;
import org.apache.asterix.lang.common.expression.QuantifiedExpression;
import org.apache.asterix.lang.common.expression.RecordConstructor;
import org.apache.asterix.lang.common.expression.UnaryExpr;
import org.apache.asterix.lang.common.expression.VariableExpr;
import org.apache.asterix.lang.common.rewrites.LangRewritingContext;
import org.apache.asterix.lang.common.rewrites.VariableSubstitutionEnvironment;
import org.apache.asterix.lang.common.statement.FunctionDecl;
import org.apache.asterix.lang.common.statement.Query;
import org.apache.asterix.lang.common.struct.Identifier;
import org.apache.asterix.lang.common.struct.QuantifiedPair;
import org.apache.asterix.lang.common.struct.VarIdentifier;
import org.apache.asterix.lang.common.util.VariableCloneAndSubstitutionUtil;
import org.apache.asterix.lang.common.visitor.base.AbstractQueryExpressionVisitor;
import org.apache.hyracks.algebricks.common.utils.Pair;

public class CloneAndSubstituteVariablesVisitor extends
        AbstractQueryExpressionVisitor<Pair<ILangExpression, VariableSubstitutionEnvironment>, VariableSubstitutionEnvironment> {

    private LangRewritingContext context;

    public CloneAndSubstituteVariablesVisitor(LangRewritingContext context) {
        this.context = context;
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(LetClause lc,
            VariableSubstitutionEnvironment env) throws CompilationException {
        Pair<ILangExpression, VariableSubstitutionEnvironment> p1 = lc.getBindingExpr().accept(this, env);
        VariableExpr varExpr = lc.getVarExpr();
        VariableExpr newVe = generateNewVariable(context, varExpr);
        LetClause newLet = new LetClause(newVe, (Expression) p1.first);
        newLet.setSourceLocation(lc.getSourceLocation());
        return new Pair<>(newLet, VariableCloneAndSubstitutionUtil.eliminateSubstFromList(lc.getVarExpr(), env));
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(GroupbyClause gc,
            VariableSubstitutionEnvironment env) throws CompilationException {
        VariableSubstitutionEnvironment newSubs = env;
        List<List<GbyVariableExpressionPair>> gbyList = gc.getGbyPairList();
        List<List<GbyVariableExpressionPair>> newGbyList = new ArrayList<>(gbyList.size());
        for (List<GbyVariableExpressionPair> gbyPairList : gbyList) {
            newGbyList.add(VariableCloneAndSubstitutionUtil.substInVarExprPair(context, gbyPairList, newSubs, this));
        }
        List<GbyVariableExpressionPair> newDecorList = gc.hasDecorList()
                ? VariableCloneAndSubstitutionUtil.substInVarExprPair(context, gc.getDecorPairList(), newSubs, this)
                : new ArrayList<>();

        VariableExpr newGroupVar = null;
        if (gc.hasGroupVar()) {
            newGroupVar = generateNewVariable(context, gc.getGroupVar());
        }
        Map<Expression, VariableExpr> newWithMap = new HashMap<>();
        if (gc.hasWithMap()) {
            for (Entry<Expression, VariableExpr> entry : gc.getWithVarMap().entrySet()) {
                Expression newKeyVar = (Expression) entry.getKey().accept(this, env).first;
                VariableExpr newValueVar = generateNewVariable(context, entry.getValue());
                newWithMap.put(newKeyVar, newValueVar);
            }
        }
        List<Pair<Expression, Identifier>> newGroupFieldList = gc.hasGroupFieldList()
                ? VariableCloneAndSubstitutionUtil.substInFieldList(gc.getGroupFieldList(), env, this) : null;
        GroupbyClause newGroup = new GroupbyClause(newGbyList, newDecorList, newWithMap, newGroupVar, newGroupFieldList,
                gc.hasHashGroupByHint(), gc.isGroupAll());
        newGroup.setSourceLocation(gc.getSourceLocation());
        return new Pair<>(newGroup, newSubs);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(QuantifiedExpression qe,
            VariableSubstitutionEnvironment env) throws CompilationException {
        List<QuantifiedPair> oldPairs = qe.getQuantifiedList();
        List<QuantifiedPair> newPairs = new ArrayList<>(oldPairs.size());
        VariableSubstitutionEnvironment newSubs = env;
        for (QuantifiedPair t : oldPairs) {
            VariableExpr newVar = generateNewVariable(context, t.getVarExpr());
            newSubs = VariableCloneAndSubstitutionUtil.eliminateSubstFromList(newVar, newSubs);
            Pair<ILangExpression, VariableSubstitutionEnvironment> p1 =
                    visitUnnestBindingExpression(t.getExpr(), newSubs);
            QuantifiedPair t2 = new QuantifiedPair(newVar, (Expression) p1.first);
            newPairs.add(t2);
        }
        Pair<ILangExpression, VariableSubstitutionEnvironment> p2 = qe.getSatisfiesExpr().accept(this, newSubs);
        QuantifiedExpression qe2 = new QuantifiedExpression(qe.getQuantifier(), newPairs, (Expression) p2.first);
        qe2.setSourceLocation(qe.getSourceLocation());
        qe2.addHints(qe.getHints());
        return new Pair<>(qe2, newSubs);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(WhereClause wc,
            VariableSubstitutionEnvironment env) throws CompilationException {
        Pair<ILangExpression, VariableSubstitutionEnvironment> p1 = wc.getWhereExpr().accept(this, env);
        WhereClause newW = new WhereClause((Expression) p1.first);
        newW.setSourceLocation(wc.getSourceLocation());
        return new Pair<>(newW, p1.second);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(CallExpr callExpr,
            VariableSubstitutionEnvironment env) throws CompilationException {
        List<Expression> exprList =
                VariableCloneAndSubstitutionUtil.visitAndCloneExprList(callExpr.getExprList(), env, this);
        Expression newFilterExpr = null;
        if (callExpr.hasAggregateFilterExpr()) {
            Pair<ILangExpression, VariableSubstitutionEnvironment> paf =
                    callExpr.getAggregateFilterExpr().accept(this, env);
            newFilterExpr = (Expression) paf.first;
        }
        CallExpr f = new CallExpr(callExpr.getFunctionSignature(), exprList, newFilterExpr);
        f.setSourceLocation(callExpr.getSourceLocation());
        f.addHints(callExpr.getHints());
        return new Pair<>(f, env);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(FunctionDecl fd,
            VariableSubstitutionEnvironment env) throws CompilationException {
        List<VarIdentifier> newList = new ArrayList<>(fd.getParamList().size());
        for (VarIdentifier vi : fd.getParamList()) {
            VariableExpr varExpr = new VariableExpr(vi);
            if (!env.constainsOldVar(varExpr)) {
                throw new CompilationException(ErrorCode.COMPILATION_ERROR, fd.getSourceLocation(),
                        "Parameter " + vi + " does not appear in the substitution list.");
            }
            Expression newExpr = env.findSubstitution(varExpr);
            if (newExpr.getKind() != Kind.VARIABLE_EXPRESSION) {
                throw new CompilationException(ErrorCode.COMPILATION_ERROR, fd.getSourceLocation(),
                        "Parameter " + vi + " cannot be substituted by a non-variable expression.");
            }
            newList.add(((VariableExpr) newExpr).getVar());
        }

        Pair<ILangExpression, VariableSubstitutionEnvironment> p1 = fd.getFuncBody().accept(this, env);
        FunctionDecl newF = new FunctionDecl(fd.getSignature(), newList, (Expression) p1.first, fd.isStored());
        newF.setSourceLocation(fd.getSourceLocation());
        return new Pair<>(newF, env);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(IfExpr ifexpr,
            VariableSubstitutionEnvironment env) throws CompilationException {
        Pair<ILangExpression, VariableSubstitutionEnvironment> p1 = ifexpr.getCondExpr().accept(this, env);
        Pair<ILangExpression, VariableSubstitutionEnvironment> p2 = ifexpr.getThenExpr().accept(this, env);
        Pair<ILangExpression, VariableSubstitutionEnvironment> p3 = ifexpr.getElseExpr().accept(this, env);
        IfExpr i = new IfExpr((Expression) p1.first, (Expression) p2.first, (Expression) p3.first);
        i.setSourceLocation(ifexpr.getSourceLocation());
        i.addHints(ifexpr.getHints());
        return new Pair<>(i, env);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(LimitClause lc,
            VariableSubstitutionEnvironment env) throws CompilationException {
        Expression newLimitExpr = null;
        if (lc.hasLimitExpr()) {
            Pair<ILangExpression, VariableSubstitutionEnvironment> p1 = lc.getLimitExpr().accept(this, env);
            newLimitExpr = (Expression) p1.first;
        }
        Expression newOffsetExpr = null;
        if (lc.hasOffset()) {
            Pair<ILangExpression, VariableSubstitutionEnvironment> p2 = lc.getOffset().accept(this, env);
            newOffsetExpr = (Expression) p2.first;
        }
        LimitClause c = new LimitClause(newLimitExpr, newOffsetExpr);
        c.setSourceLocation(lc.getSourceLocation());
        return new Pair<>(c, env);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(ListConstructor lc,
            VariableSubstitutionEnvironment env) throws CompilationException {
        List<Expression> oldExprList = lc.getExprList();
        List<Expression> exprs = VariableCloneAndSubstitutionUtil.visitAndCloneExprList(oldExprList, env, this);
        ListConstructor c = new ListConstructor(lc.getType(), exprs);
        c.setSourceLocation(lc.getSourceLocation());
        c.addHints(lc.getHints());
        return new Pair<>(c, env);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(LiteralExpr l,
            VariableSubstitutionEnvironment env) throws CompilationException {
        return new Pair<>(l, env);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(OperatorExpr op,
            VariableSubstitutionEnvironment env) throws CompilationException {
        List<Expression> oldExprList = op.getExprList();
        List<Expression> exprs = new ArrayList<>(oldExprList.size());
        for (Expression e : oldExprList) {
            Pair<ILangExpression, VariableSubstitutionEnvironment> p1 = e.accept(this, env);
            exprs.add((Expression) p1.first);
        }
        OperatorExpr oe = new OperatorExpr(exprs, op.getOpList(), op.isCurrentop());
        oe.setSourceLocation(op.getSourceLocation());
        oe.addHints(op.getHints());
        return new Pair<>(oe, env);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(OrderbyClause oc,
            VariableSubstitutionEnvironment env) throws CompilationException {
        List<Expression> exprList =
                VariableCloneAndSubstitutionUtil.visitAndCloneExprList(oc.getOrderbyList(), env, this);
        OrderbyClause oc2 = new OrderbyClause(exprList, new ArrayList<>(oc.getModifierList()),
                new ArrayList<>(oc.getNullModifierList()));
        oc2.setNumFrames(oc.getNumFrames());
        oc2.setNumTuples(oc.getNumTuples());
        oc2.setRangeMap(oc.getRangeMap());
        oc2.setSourceLocation(oc.getSourceLocation());
        return new Pair<>(oc2, env);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(Query q, VariableSubstitutionEnvironment env)
            throws CompilationException {
        Query newQ = new Query(q.isExplain());
        Pair<ILangExpression, VariableSubstitutionEnvironment> p1 = q.getBody().accept(this, env);
        newQ.setBody((Expression) p1.first);
        newQ.setSourceLocation(q.getSourceLocation());
        return new Pair<>(newQ, p1.second);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(RecordConstructor rc,
            VariableSubstitutionEnvironment env) throws CompilationException {
        List<FieldBinding> oldFbs = rc.getFbList();
        ArrayList<FieldBinding> newFbs = new ArrayList<>(oldFbs.size());
        for (FieldBinding fb : oldFbs) {
            Pair<ILangExpression, VariableSubstitutionEnvironment> p1 = fb.getLeftExpr().accept(this, env);
            Pair<ILangExpression, VariableSubstitutionEnvironment> p2 = fb.getRightExpr().accept(this, env);
            FieldBinding fb2 = new FieldBinding((Expression) p1.first, (Expression) p2.first);
            newFbs.add(fb2);
        }
        RecordConstructor newRc = new RecordConstructor(newFbs);
        newRc.setSourceLocation(rc.getSourceLocation());
        newRc.addHints(rc.getHints());
        return new Pair<>(newRc, env);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(UnaryExpr u,
            VariableSubstitutionEnvironment env) throws CompilationException {
        Pair<ILangExpression, VariableSubstitutionEnvironment> p1 = u.getExpr().accept(this, env);
        UnaryExpr newU = new UnaryExpr(u.getExprType(), (Expression) p1.first);
        newU.setSourceLocation(u.getSourceLocation());
        newU.addHints(u.getHints());
        return new Pair<>(newU, env);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(IndexAccessor ia,
            VariableSubstitutionEnvironment env) throws CompilationException {
        Pair<ILangExpression, VariableSubstitutionEnvironment> p1 = ia.getExpr().accept(this, env);
        Expression indexExpr = null;
        if (ia.getIndexExpr() != null) {
            Pair<ILangExpression, VariableSubstitutionEnvironment> p2 = ia.getIndexExpr().accept(this, env);
            indexExpr = (Expression) p2.first;
        }
        IndexAccessor i = new IndexAccessor((Expression) p1.first, ia.getIndexKind(), indexExpr);
        i.setSourceLocation(ia.getSourceLocation());
        i.addHints(ia.getHints());
        return new Pair<>(i, env);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(ListSliceExpression expression,
            VariableSubstitutionEnvironment env) throws CompilationException {
        Pair<ILangExpression, VariableSubstitutionEnvironment> expressionPair = expression.getExpr().accept(this, env);
        Expression startIndexExpression;
        Expression endIndexExpression = null;

        // Start index expression
        Pair<ILangExpression, VariableSubstitutionEnvironment> startExpressionPair =
                expression.getStartIndexExpression().accept(this, env);
        startIndexExpression = (Expression) startExpressionPair.first;

        // End index expression (optional)
        if (expression.hasEndExpression()) {
            Pair<ILangExpression, VariableSubstitutionEnvironment> endExpressionPair =
                    expression.getEndIndexExpression().accept(this, env);
            endIndexExpression = (Expression) endExpressionPair.first;
        }

        // Resulted expression
        ListSliceExpression resultExpression =
                new ListSliceExpression((Expression) expressionPair.first, startIndexExpression, endIndexExpression);
        resultExpression.setSourceLocation(expression.getSourceLocation());
        resultExpression.addHints(expression.getHints());
        return new Pair<>(resultExpression, env);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(IVisitorExtension ve,
            VariableSubstitutionEnvironment arg) throws CompilationException {
        return ve.remapCloneDispatch(this, arg);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(FieldAccessor fa,
            VariableSubstitutionEnvironment env) throws CompilationException {
        Pair<ILangExpression, VariableSubstitutionEnvironment> p = fa.getExpr().accept(this, env);
        FieldAccessor newF = new FieldAccessor((Expression) p.first, fa.getIdent());
        newF.setSourceLocation(fa.getSourceLocation());
        newF.addHints(fa.getHints());
        return new Pair<>(newF, p.second);
    }

    @Override
    public Pair<ILangExpression, VariableSubstitutionEnvironment> visit(VariableExpr v,
            VariableSubstitutionEnvironment env) throws CompilationException {
        return new Pair<>(rewriteVariableExpr(v, env), env);
    }

    // Replace a variable expression if the variable is to-be substituted.
    protected Expression rewriteVariableExpr(VariableExpr expr, VariableSubstitutionEnvironment env)
            throws CompilationException {
        if (env.constainsOldVar(expr)) {
            return env.findSubstitution(expr);
        } else {
            // it is a variable from the context
            VarIdentifier var = context.getRewrittenVar(expr.getVar().getId());
            if (var != null) {
                VariableExpr newVarExpr = new VariableExpr(var);
                newVarExpr.setSourceLocation(expr.getSourceLocation());
                newVarExpr.addHints(expr.getHints());
                return newVarExpr;
            }
        }
        return expr;
    }

    /**
     * Generates a new variable for an existing variable.
     *
     * @param context
     *            , the language rewriting context which keeps all the rewriting variable-int-id to variable-string-identifier mappings.
     * @param varExpr
     *            , the existing variable expression.
     * @return the new variable expression.
     */
    public VariableExpr generateNewVariable(LangRewritingContext context, VariableExpr varExpr) {
        VarIdentifier vi = varExpr.getVar();
        VarIdentifier newVar = context.mapOldId(vi.getId(), vi.getValue());
        VariableExpr newVarExpr = new VariableExpr(newVar);
        newVarExpr.setSourceLocation(varExpr.getSourceLocation());
        return newVarExpr;
    }

    /**
     * Visits an expression that is used for unnest binding.
     *
     * @param expr,
     *            the expression to consider.
     * @param env,
     *            the variable substitution environment.
     * @return a pair of an ILangExpression and a variable substitution environment.
     * @throws CompilationException
     */
    protected Pair<ILangExpression, VariableSubstitutionEnvironment> visitUnnestBindingExpression(Expression expr,
            VariableSubstitutionEnvironment env) throws CompilationException {
        return expr.accept(this, env);
    }

}
