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.CountDownLatch;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import org.junit.Assert;
import org.junit.Test;

/**
 * Engine 级有界分片调度器测试。
 */
public class FederationFragmentSchedulerTest {

    /**
     * 验证两个分片任务能够真实重叠执行，且调度器关闭后拒绝新任务。
     *
     * @throws Exception 等待任务失败
     */
    @Test
    public void shouldExecuteFragmentTasksConcurrentlyAndRejectAfterClose() throws Exception {
        FederationFragmentScheduler scheduler = new FederationFragmentScheduler(
            FederationExecutionPolicy.basic()
        );
        CountDownLatch started = new CountDownLatch(2);
        CountDownLatch release = new CountDownLatch(1);
        AtomicInteger active = new AtomicInteger();
        AtomicInteger maximumActive = new AtomicInteger();
        ExecutorCompletionService<Void> completions = new ExecutorCompletionService<>(scheduler);
        try {
            Future<Void> first = completions.submit(
                () -> runBlockingTask(started, release, active, maximumActive)
            );
            Future<Void> second = completions.submit(
                () -> runBlockingTask(started, release, active, maximumActive)
            );

            Assert.assertTrue(started.await(2, TimeUnit.SECONDS));
            Assert.assertEquals(2, maximumActive.get());
            release.countDown();
            first.get(2, TimeUnit.SECONDS);
            second.get(2, TimeUnit.SECONDS);
        } finally {
            release.countDown();
            scheduler.close();
        }

        try {
            scheduler.execute(() -> { });
            Assert.fail("closed scheduler should reject new tasks");
        } catch (FederationSqlException exception) {
            Assert.assertEquals(FederationSqlErrorCode.ENGINE_CLOSED, exception.errorCode());
        }
    }

    /**
     * 验证队列饱和时明确拒绝，且不会在提交线程绕过 Worker 上限执行任务。
     *
     * @throws Exception 等待任务失败
     */
    @Test
    public void shouldRejectSaturationWithoutRunningTaskOnCaller() throws Exception {
        FederationFragmentScheduler scheduler = new FederationFragmentScheduler(1, 1);
        CountDownLatch firstStarted = new CountDownLatch(1);
        CountDownLatch releaseFirst = new CountDownLatch(1);
        CountDownLatch queuedCompleted = new CountDownLatch(1);
        AtomicBoolean rejectedTaskRan = new AtomicBoolean();
        try {
            scheduler.execute(() -> {
                firstStarted.countDown();
                try {
                    releaseFirst.await(2, TimeUnit.SECONDS);
                } catch (InterruptedException exception) {
                    Thread.currentThread().interrupt();
                }
            });
            Assert.assertTrue(firstStarted.await(2, TimeUnit.SECONDS));
            scheduler.execute(queuedCompleted::countDown);

            try {
                scheduler.execute(() -> rejectedTaskRan.set(true));
                Assert.fail("saturated scheduler should reject the task");
            } catch (FederationSqlException exception) {
                Assert.assertEquals(
                    FederationSqlErrorCode.QUERY_ADMISSION_TIMEOUT,
                    exception.errorCode()
                );
            }
            Assert.assertFalse(rejectedTaskRan.get());
        } finally {
            releaseFirst.countDown();
            Assert.assertTrue(queuedCompleted.await(2, TimeUnit.SECONDS));
            scheduler.close();
        }
    }

    private static Void runBlockingTask(
        CountDownLatch started,
        CountDownLatch release,
        AtomicInteger active,
        AtomicInteger maximumActive
    ) throws InterruptedException {
        int current = active.incrementAndGet();
        maximumActive.accumulateAndGet(current, Math::max);
        started.countDown();
        try {
            Assert.assertTrue(release.await(2, TimeUnit.SECONDS));
            return null;
        } finally {
            active.decrementAndGet();
        }
    }
}
