package com.easyagents.federation.sql.execute;

import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import com.easyagents.federation.sql.source.SourceId;
import java.time.Duration;
import java.util.function.BooleanSupplier;

/**
 * 可替换的查询并发准入控制器。
 */
@FunctionalInterface
public interface FederationQueryAdmissionController extends AutoCloseable {

    /**
     * 获取查询许可。
     *
     * @param sourceId 数据源标识
     * @param queryId 查询标识
     * @param timeout 最大等待时间
     * @return 查询许可
     */
    FederationQueryPermit acquire(SourceId sourceId, QueryId queryId, Duration timeout);

    /**
     * 获取支持查询级取消的许可。
     *
     * <p>自定义实现可以覆盖此方法及时中断分布式或远程准入等待。</p>
     *
     * @param sourceId 数据源标识
     * @param queryId 查询标识
     * @param timeout 最大等待时间
     * @param cancellationRequested 取消状态
     * @return 查询许可
     */
    default FederationQueryPermit acquire(
        SourceId sourceId,
        QueryId queryId,
        Duration timeout,
        BooleanSupplier cancellationRequested
    ) {
        return acquire(sourceId, queryId, timeout);
    }

    /**
     * 为一次单源或联邦查询获取一份查询级许可。
     *
     * <p>兼容实现只接受单源请求。联邦查询必须由实现方明确覆盖本方法，避免其余
     * 物理源静默绕过源级配额。</p>
     *
     * @param request 查询级准入请求
     * @return 查询许可
     */
    default FederationQueryPermit acquire(QueryAdmissionRequest request) {
        if (request.sourceIds().size() != 1) {
            throw new FederationSqlException(
                FederationSqlErrorCode.INVALID_QUERY_SCOPE,
                "admission controller does not declare multi-source query support"
            );
        }
        return acquire(
            request.primarySourceId(),
            request.queryId(),
            request.timeout(),
            request.cancellationRequested()
        );
    }

    /**
     * 关闭控制器；默认无额外资源。
     */
    @Override
    default void close() {
    }

    /**
     * 返回无并发限制的控制器。
     *
     * @return 无限制控制器
     */
    static FederationQueryAdmissionController unlimited() {
        return new FederationQueryAdmissionController() {
            @Override
            public FederationQueryPermit acquire(
                SourceId sourceId,
                QueryId queryId,
                Duration timeout
            ) {
                return () -> { };
            }

            @Override
            public FederationQueryPermit acquire(QueryAdmissionRequest request) {
                return () -> { };
            }
        };
    }
}
