package com.easyagents.federation.sql.runtime;

import org.apache.calcite.DataContext;
import org.apache.calcite.linq4j.AbstractEnumerable;
import org.apache.calcite.linq4j.Enumerable;
import org.apache.calcite.linq4j.Enumerator;

/**
 * 为生成的 Calcite Enumerable 提供运行期本地算子预算检查。
 */
public final class FederationBudgetEnumerable {

    /**
     * 禁止实例化运行期桥接工具。
     */
    private FederationBudgetEnumerable() {
    }

    /**
     * 包装一个本地算子输出，在每行交给下游前计量资源并检查取消与时限。
     *
     * @param input Calcite 算子输出
     * @param context 当前查询 DataContext
     * @param operatorName 本地算子名称
     * @param <T> Calcite 内部行类型
     * @return 带预算检查的 Enumerable
     */
    public static <T> Enumerable<T> wrap(
        Enumerable<T> input,
        DataContext context,
        String operatorName
    ) {
        Object value = context.get(FederationExecutionSession.DATA_CONTEXT_KEY);
        if (!(value instanceof FederationExecutionSession session)) {
            throw new IllegalStateException("federation execution session is missing");
        }
        return new AbstractEnumerable<>() {
            @Override
            public Enumerator<T> enumerator() {
                Enumerator<T> delegate = input.enumerator();
                return new Enumerator<>() {
                    @Override
                    public T current() {
                        return delegate.current();
                    }

                    @Override
                    public boolean moveNext() {
                        session.ensureExecutionAllowed();
                        long started = System.nanoTime();
                        if (!delegate.moveNext()) {
                            return false;
                        }
                        session.recordLocalIntermediate(
                            operatorName,
                            delegate.current(),
                            System.nanoTime() - started
                        );
                        return true;
                    }

                    @Override
                    public void reset() {
                        delegate.reset();
                    }

                    @Override
                    public void close() {
                        delegate.close();
                    }
                };
            }
        };
    }
}
