package com.easyagents.federation.sql.runtime;

import java.util.List;
import org.apache.calcite.DataContext;
import org.apache.calcite.adapter.enumerable.EnumerableConvention;
import org.apache.calcite.adapter.enumerable.EnumerableRel;
import org.apache.calcite.adapter.enumerable.EnumerableRelImplementor;
import org.apache.calcite.linq4j.tree.BlockBuilder;
import org.apache.calcite.linq4j.tree.Expression;
import org.apache.calcite.linq4j.tree.Expressions;
import org.apache.calcite.plan.RelTraitSet;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.SingleRel;

/**
 * 在 Calcite Enumerable 本地算子输出边界插入查询预算检查。
 */
final class EnumerableFederationBudgetRel extends SingleRel implements EnumerableRel {

    private final String operatorName;

    /**
     * 创建预算边界。
     *
     * @param input 已转换为 Enumerable Convention 的输入
     * @param operatorName 被计量的本地算子名称
     */
    EnumerableFederationBudgetRel(RelNode input, String operatorName) {
        super(input.getCluster(), input.getTraitSet(), input);
        if (input.getConvention() != EnumerableConvention.INSTANCE) {
            throw new IllegalArgumentException("budget input must use Enumerable convention");
        }
        this.operatorName = operatorName;
    }

    /** {@inheritDoc} */
    @Override
    public EnumerableFederationBudgetRel copy(
        RelTraitSet traitSet,
        List<RelNode> inputs
    ) {
        return new EnumerableFederationBudgetRel(sole(inputs), operatorName);
    }

    /** {@inheritDoc} */
    @Override
    public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
        BlockBuilder builder = new BlockBuilder();
        Result child = implementor.visitChild(this, 0, (EnumerableRel) getInput(), pref);
        Expression enumerable = builder.append("budgetInput", child.block);
        Expression guarded = Expressions.call(
            FederationBudgetEnumerable.class,
            "wrap",
            enumerable,
            DataContext.ROOT,
            Expressions.constant(operatorName)
        );
        builder.add(Expressions.return_(null, guarded));
        return implementor.result(child.physType, builder.toBlock());
    }
}
