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.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.BooleanSupplier;

/**
 * 使用公平信号量限制单节点查询并发的默认控制器。
 */
public final class LocalFederationQueryAdmissionController implements FederationQueryAdmissionController {

    private final int maxConcurrentQueries;
    private final Semaphore permits;
    private final AtomicBoolean closed = new AtomicBoolean();

    /**
     * 创建本地准入控制器。
     *
     * @param maxConcurrentQueries 最大并发查询数
     */
    public LocalFederationQueryAdmissionController(int maxConcurrentQueries) {
        if (maxConcurrentQueries <= 0) {
            throw new IllegalArgumentException("maxConcurrentQueries must be positive");
        }
        this.maxConcurrentQueries = maxConcurrentQueries;
        this.permits = new Semaphore(maxConcurrentQueries, true);
    }

    /**
     * 在指定时间内获取本地许可。
     *
     * @param sourceId 数据源标识
     * @param queryId 查询标识
     * @param timeout 最大等待时间
     * @return 可幂等关闭的许可
     */
    @Override
    public FederationQueryPermit acquire(SourceId sourceId, QueryId queryId, Duration timeout) {
        return acquire(sourceId, queryId, timeout, () -> false);
    }

    /**
     * 在等待本地许可期间轮询查询取消状态。
     *
     * @param sourceId 数据源标识
     * @param queryId 查询标识
     * @param timeout 最大等待时间
     * @param cancellationRequested 取消状态
     * @return 可幂等关闭的许可
     */
    @Override
    public FederationQueryPermit acquire(
        SourceId sourceId,
        QueryId queryId,
        Duration timeout,
        BooleanSupplier cancellationRequested
    ) {
        ensureOpen();
        if (cancellationRequested.getAsBoolean()) {
            throw cancelled(queryId);
        }
        boolean acquired = false;
        try {
            long timeoutNanos = timeout.toNanos();
            if (timeoutNanos == 0) {
                acquired = permits.tryAcquire();
            } else {
                long started = System.nanoTime();
                long remaining = timeoutNanos;
                long pollNanos = TimeUnit.MILLISECONDS.toNanos(50);
                while (!acquired && remaining > 0) {
                    acquired = permits.tryAcquire(Math.min(remaining, pollNanos), TimeUnit.NANOSECONDS);
                    if (!acquired && cancellationRequested.getAsBoolean()) {
                        throw cancelled(queryId);
                    }
                    remaining = timeoutNanos - (System.nanoTime() - started);
                }
            }
        } catch (InterruptedException exception) {
            Thread.currentThread().interrupt();
            throw new FederationSqlException(
                FederationSqlErrorCode.QUERY_ADMISSION_TIMEOUT,
                "query admission was interrupted for source " + sourceId,
                exception
            );
        }
        if (cancellationRequested.getAsBoolean()) {
            if (acquired) {
                permits.release();
            }
            throw cancelled(queryId);
        }
        if (!acquired) {
            throw new FederationSqlException(
                FederationSqlErrorCode.QUERY_ADMISSION_TIMEOUT,
                "query admission timed out for source " + sourceId
            );
        }
        if (closed.get()) {
            permits.release();
            throw new FederationSqlException(
                FederationSqlErrorCode.ENGINE_CLOSED,
                "query admission controller is closed"
            );
        }
        AtomicBoolean released = new AtomicBoolean();
        return () -> {
            if (released.compareAndSet(false, true)) {
                permits.release();
            }
        };
    }

    /**
     * 为单源或联邦查询获取一份节点级本地许可。
     *
     * @param request 查询级准入请求
     * @return 可幂等关闭的节点许可
     */
    @Override
    public FederationQueryPermit acquire(QueryAdmissionRequest request) {
        return acquire(
            request.primarySourceId(),
            request.queryId(),
            request.timeout(),
            request.cancellationRequested()
        );
    }

    /**
     * 关闭控制器并唤醒等待准入的线程。
     */
    @Override
    public void close() {
        if (closed.compareAndSet(false, true)) {
            permits.release(maxConcurrentQueries);
        }
    }

    private void ensureOpen() {
        if (closed.get()) {
            throw new FederationSqlException(
                FederationSqlErrorCode.ENGINE_CLOSED,
                "query admission controller is closed"
            );
        }
    }

    private static FederationSqlException cancelled(QueryId queryId) {
        return new FederationSqlException(
            FederationSqlErrorCode.QUERY_CANCELLED,
            "query admission was cancelled: " + queryId.value()
        );
    }
}
