/*
 * 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.pig.backend.hadoop.executionengine.mapReduceLayer;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import org.apache.pig.FuncSpec;
import org.apache.pig.PigException;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.plans.MROpPlanVisitor;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.plans.MROperPlan;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.PhysicalOperator;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POLimit;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POLoad;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POLocalRearrange;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POPackage;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POStore;
import org.apache.pig.impl.PigContext;
import org.apache.pig.impl.io.FileLocalizer;
import org.apache.pig.impl.io.FileSpec;
import org.apache.pig.impl.plan.DepthFirstWalker;
import org.apache.pig.impl.plan.NodeIdGenerator;
import org.apache.pig.impl.plan.OperatorKey;
import org.apache.pig.impl.plan.PlanException;
import org.apache.pig.impl.plan.VisitorException;
import org.apache.pig.impl.util.Utils;

public class LimitAdjuster extends MROpPlanVisitor {
    ArrayList<MapReduceOper> opsToAdjust = new ArrayList<MapReduceOper>();
    PigContext pigContext;
    NodeIdGenerator nig;
    private String scope;


    public LimitAdjuster(MROperPlan plan, PigContext pigContext) {
        super(plan, new DepthFirstWalker<MapReduceOper, MROperPlan>(plan));
        this.pigContext = pigContext;
        nig = NodeIdGenerator.getGenerator();
        List<MapReduceOper> roots = plan.getRoots();
        scope = roots.get(0).getOperatorKey().getScope();
    }

    @Override
    public void visitMROp(MapReduceOper mr) throws VisitorException {
        // Look for map reduce operators which contains limit operator.
        // If so, add one additional map-reduce
        // operator with 1 reducer into the original plan.

        // TODO: This new MR job can be skipped if at runtime we discover that
        // its parent only has a single reducer (mr.requestedParallelism!=1).
        // This check MUST happen at runtime since that's when reducer estimation happens.
        if ((mr.limit!=-1 || mr.limitPlan!=null) )
        {
            opsToAdjust.add(mr);
        }
    }

    public void adjust() throws IOException, PlanException
    {
        for (MapReduceOper mr:opsToAdjust)
        {
            if (mr.reducePlan.isEmpty()) continue;
            List<PhysicalOperator> mpLeaves = mr.reducePlan.getLeaves();
            if (mpLeaves.size() != 1) {
                int errCode = 2024;
                String msg = "Expected reduce to have single leaf. Found " + mpLeaves.size() + " leaves.";
                throw new MRCompilerException(msg, errCode, PigException.BUG);
            }
            PhysicalOperator mpLeaf = mpLeaves.get(0);
            if (!pigContext.inIllustrator) {
                if (!(mpLeaf instanceof POStore)) {
                    int errCode = 2025;
                    String msg = "Expected leaf of reduce plan to " +
                        "always be POStore. Found " + mpLeaf.getClass().getSimpleName();
                    throw new MRCompilerException(msg, errCode, PigException.BUG);
                }
            }
            FileSpec oldSpec = ((POStore)mpLeaf).getSFile();
            boolean oldIsTmpStore = ((POStore)mpLeaf).isTmpStore();

            FileSpec fSpec = new FileSpec(FileLocalizer.getTemporaryPath(pigContext).toString(),
                    new FuncSpec(Utils.getTmpFileCompressorName(pigContext)));
            POStore storeOp = (POStore) mpLeaf;
            storeOp.setSFile(fSpec);
            storeOp.setIsTmpStore(true);
            mr.setReduceDone(true);
            MapReduceOper limitAdjustMROp = new MapReduceOper(new OperatorKey(scope,nig.getNextNodeId(scope)));
            POLoad ld = new POLoad(new OperatorKey(scope,nig.getNextNodeId(scope)));
            ld.setPc(pigContext);
            ld.setLFile(fSpec);
            ld.setIsTmpLoad(true);
            limitAdjustMROp.mapPlan.add(ld);
            if (mr.isGlobalSort()) {
                connectMapToReduceLimitedSort(limitAdjustMROp, mr);
            } else {
                MRUtil.simpleConnectMapToReduce(limitAdjustMROp, scope, nig);
            }
            // Need to split the original reduce plan into two mapreduce job:
            // 1st: From the root(POPackage) to POLimit
            // 2nd: From POLimit to leaves(POStore), duplicate POLimit
            // The reason for doing that:
            // 1. We need to have two map-reduce job, otherwise, we will end up with
            //    N*M records, N is number of reducer, M is limit constant. We need
            //    one extra mapreduce job with 1 reducer
            // 2. We don't want to move operator after POLimit into the first mapreduce
            //    job, because:
            //    * Foreach will shift the key type for second mapreduce job, see PIG-461
            //    * Foreach flatten may generating more than M records, which get cut
            //      by POLimit, see PIG-2231
            splitReducerForLimit(limitAdjustMROp, mr);

            if (mr.isGlobalSort())
            {
                limitAdjustMROp.setLimitAfterSort(true);
                limitAdjustMROp.setSortOrder(mr.getSortOrder());
            }

            POStore st = new POStore(new OperatorKey(scope,nig.getNextNodeId(scope)));
            st.setSFile(oldSpec);
            st.setIsTmpStore(oldIsTmpStore);
            st.setSchema(((POStore)mpLeaf).getSchema());
            st.setSignature(((POStore)mpLeaf).getSignature());
            st.copyAliasFrom(mpLeaf);

            limitAdjustMROp.reducePlan.addAsLeaf(st);
            limitAdjustMROp.requestedParallelism = 1;
            limitAdjustMROp.setLimitOnly(true);

            List<MapReduceOper> successorList = mPlan.getSuccessors(mr);
            MapReduceOper successors[] = null;

            // Save a snapshot for successors, since we will modify MRPlan,
            // use the list directly will be problematic
            if (successorList!=null && successorList.size()>0)
            {
                successors = new MapReduceOper[successorList.size()];
                int i=0;
                for (MapReduceOper op:successorList)
                    successors[i++] = op;
            }

            // Process UDFs
            for (String udf : mr.UDFs) {
                if (!limitAdjustMROp.UDFs.contains(udf)) {
                    limitAdjustMROp.UDFs.add(udf);
                }
            }

            mPlan.add(limitAdjustMROp);
            mPlan.connect(mr, limitAdjustMROp);

            if (successors!=null)
            {
                for (int i=0;i<successors.length;i++)
                {
                    MapReduceOper nextMr = successors[i];
                    if (nextMr!=null)
                        mPlan.disconnect(mr, nextMr);

                    if (nextMr!=null)
                        mPlan.connect(limitAdjustMROp, nextMr);
                }
            }
        }
    }

    // Move all operators between POLimit and POStore in reducer plan
    // from firstMROp to the secondMROp
    private void splitReducerForLimit(MapReduceOper secondMROp,
            MapReduceOper firstMROp) throws PlanException, VisitorException {

        PhysicalOperator op = firstMROp.reducePlan.getRoots().get(0);
        assert(op instanceof POPackage);

        while (true) {
            List<PhysicalOperator> succs = firstMROp.reducePlan
                    .getSuccessors(op);
            if (succs==null) break;
            op = succs.get(0);
            if (op instanceof POLimit) {
                // find operator after POLimit
                op = firstMROp.reducePlan.getSuccessors(op).get(0);
                break;
            }
        }

        POLimit pLimit2 = new POLimit(new OperatorKey(scope,nig.getNextNodeId(scope)));
        pLimit2.setLimit(firstMROp.limit);
        pLimit2.setLimitPlan(firstMROp.limitPlan);
        secondMROp.reducePlan.addAsLeaf(pLimit2);

        while (true) {
            if (op instanceof POStore) break;
            PhysicalOperator opToMove = op;
            List<PhysicalOperator> succs = firstMROp.reducePlan
                    .getSuccessors(op);
            op = succs.get(0);

            firstMROp.reducePlan.removeAndReconnect(opToMove);
            secondMROp.reducePlan.addAsLeaf(opToMove);

        }
    }

    private void connectMapToReduceLimitedSort(MapReduceOper mro, MapReduceOper sortMROp) throws PlanException, VisitorException
    {
        POLocalRearrange slr = (POLocalRearrange)sortMROp.mapPlan.getLeaves().get(0);

        POLocalRearrange lr = null;
        try {
            lr = slr.clone();
        } catch (CloneNotSupportedException e) {
            int errCode = 2147;
            String msg = "Error cloning POLocalRearrange for limit after sort";
            throw new MRCompilerException(msg, errCode, PigException.BUG, e);
        }

        mro.mapPlan.addAsLeaf(lr);

        POPackage spkg = (POPackage)sortMROp.reducePlan.getRoots().get(0);

        POPackage pkg = null;
        try {
            pkg = spkg.clone();
        } catch (Exception e) {
            int errCode = 2148;
            String msg = "Error cloning POPackageLite for limit after sort";
            throw new MRCompilerException(msg, errCode, PigException.BUG, e);
        }
        mro.reducePlan.add(pkg);
        mro.reducePlan.addAsLeaf(MRUtil.getPlainForEachOP(scope, nig));
    }
}