/*
 * Copyright (C) 2016-2022 ActionTech.
 * License: http://www.gnu.org/licenses/gpl.html GPL version 2 or higher.
 */

package com.actiontech.dble.plan.optimizer;

import com.actiontech.dble.plan.node.JoinNode;
import com.actiontech.dble.plan.node.MergeNode;
import com.actiontech.dble.plan.node.PlanNode;
import com.actiontech.dble.plan.node.QueryNode;
import com.actiontech.dble.singleton.TraceManager;
import com.google.common.collect.ImmutableMap;

public final class LimitPusher {
    private LimitPusher() {
    }

    public static PlanNode optimize(PlanNode qtn) {
        TraceManager.TraceObject traceObject = TraceManager.threadTrace("optimize-for-limit");
        try {
            findChild(qtn);
            return qtn;
        } finally {
            TraceManager.log(ImmutableMap.of("plan-node", qtn), traceObject);
            TraceManager.finishSpan(traceObject);
        }
    }

    private static PlanNode findChild(PlanNode qtn) {
        if (qtn instanceof MergeNode) {
            // optimizer limit
            // union: push down limit to children
            // union all:push down limit to children and add distinct
            MergeNode node = (MergeNode) qtn;
            long limitFrom = node.getLimitFrom();
            long limitTo = node.getLimitTo();
            if (limitFrom != -1 && limitTo != -1) {
                for (PlanNode child : node.getChildren()) {
                    pushLimit(child, limitFrom, limitTo, node.isUnion());
                }
            }

        } else if ((qtn instanceof JoinNode) || (qtn instanceof QueryNode)) {
            for (PlanNode child : qtn.getChildren()) {
                findChild(child);
            }
        }
        return qtn;
    }

    private static void pushLimit(PlanNode node, long limitFrom, long limitTo, boolean isUnion) {
        if (isUnion) {
            node.setDistinct(true);
        }
        if (node.getLimitFrom() == -1 && node.getLimitTo() == -1) {
            node.setLimitFrom(0);
            node.setLimitTo(limitFrom + limitTo);
        }
    }

}
