package com.easyagents.federation.sql.runtime;

import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import com.easyagents.federation.sql.federation.FederationExecutionPolicy;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.Executor;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Engine 级有界分片工作线程池；查询级并发度仍由有效执行策略单独控制。
 */
final class FederationFragmentScheduler implements Executor, AutoCloseable {

    private static final int MAXIMUM_WORKERS = 32;
    private static final int MINIMUM_QUEUE_CAPACITY = 64;
    private static final int QUEUE_CAPACITY_PER_WORKER = 8;

    private final ThreadPoolExecutor executor;
    private final AtomicBoolean closed = new AtomicBoolean();

    /**
     * 创建 Engine 独占的有界分片调度器。
     *
     * @param policy Engine 联邦执行硬上限
     */
    FederationFragmentScheduler(FederationExecutionPolicy policy) {
        int availableProcessors = Math.max(1, Runtime.getRuntime().availableProcessors());
        int workers = Math.min(
            MAXIMUM_WORKERS,
            Math.max(policy.maximumConcurrentFragments(), availableProcessors)
        );
        int queueCapacity = Math.max(
            MINIMUM_QUEUE_CAPACITY,
            workers * QUEUE_CAPACITY_PER_WORKER
        );
        this.executor = executor(workers, queueCapacity);
    }

    FederationFragmentScheduler(int workers, int queueCapacity) {
        if (workers <= 0) {
            throw new IllegalArgumentException("workers must be positive");
        }
        if (queueCapacity <= 0) {
            throw new IllegalArgumentException("queueCapacity must be positive");
        }
        this.executor = executor(workers, queueCapacity);
    }

    private static ThreadPoolExecutor executor(int workers, int queueCapacity) {
        AtomicInteger threadSequence = new AtomicInteger();
        return new ThreadPoolExecutor(
            workers,
            workers,
            0L,
            TimeUnit.MILLISECONDS,
            new ArrayBlockingQueue<>(queueCapacity),
            runnable -> {
                Thread thread = new Thread(
                    runnable,
                    "easy-agents-federation-fragment-" + threadSequence.incrementAndGet()
                );
                thread.setDaemon(true);
                return thread;
            },
            new ThreadPoolExecutor.AbortPolicy()
        );
    }

    /**
     * 提交分片任务；队列饱和时明确拒绝，禁止在调用线程绕过 Worker 上限执行 JDBC。
     *
     * @param command 分片任务
     */
    @Override
    public void execute(Runnable command) {
        if (closed.get()) {
            throw engineClosed();
        }
        try {
            executor.execute(command);
        } catch (RejectedExecutionException exception) {
            if (closed.get() || executor.isShutdown()) {
                throw engineClosed();
            }
            throw schedulerOverloaded(exception);
        }
    }

    /**
     * 中断排队任务并关闭工作线程。
     */
    @Override
    public void close() {
        if (closed.compareAndSet(false, true)) {
            executor.shutdownNow();
        }
    }

    private static FederationSqlException engineClosed() {
        return new FederationSqlException(
            FederationSqlErrorCode.ENGINE_CLOSED,
            "federation fragment scheduler is closed"
        );
    }

    private static FederationSqlException schedulerOverloaded(Throwable cause) {
        return new FederationSqlException(
            FederationSqlErrorCode.QUERY_ADMISSION_TIMEOUT,
            "federation fragment scheduler is saturated",
            cause
        );
    }
}
