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.execute.QueryId;
import com.easyagents.federation.sql.execute.StatementLifecycle;
import java.lang.reflect.Proxy;
import java.sql.Statement;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import org.junit.Assert;
import org.junit.Test;

/**
 * 查询全生命周期取消登记与终态保留测试。
 */
public class QueryCancellationRegistryTest {

    /**
     * 验证 Statement 注销和 QueryId 清理后，消费线程仍能读取取消终态。
     */
    @Test
    public void shouldRetainCancellationStateForConcurrentConsumers() throws Exception {
        QueryCancellationRegistry registry = new QueryCancellationRegistry();
        QueryId queryId = new QueryId("cancelled-query");
        QueryCancellationRegistry.QueryRegistration registration = registry.begin(queryId);
        StatementLifecycle lifecycle = registration.statementLifecycle();
        AtomicInteger cancelCalls = new AtomicInteger();
        AtomicInteger closeCalls = new AtomicInteger();
        Statement statement = statement(cancelCalls, closeCalls);

        lifecycle.register(statement);
        Assert.assertTrue(registry.cancel(queryId));
        lifecycle.unregister(statement);
        registration.close();

        awaitTermination(cancelCalls, closeCalls, 1);
        Assert.assertTrue(lifecycle.cancellationRequested());
        Assert.assertFalse(registry.cancel(queryId));
        registry.close();
    }

    /**
     * 验证在 Statement 登记前发出的取消会阻止后续 JDBC 执行。
     */
    @Test
    public void shouldRejectStatementRegisteredAfterPendingCancellation() {
        QueryCancellationRegistry registry = new QueryCancellationRegistry();
        QueryId queryId = new QueryId("pending-cancel");
        QueryCancellationRegistry.QueryRegistration registration = registry.begin(queryId);
        Assert.assertTrue(registry.cancel(queryId));

        try {
            registration.statementLifecycle().register(statement(
                new AtomicInteger(),
                new AtomicInteger()
            ));
            Assert.fail("cancelled query should reject a late Statement");
        } catch (FederationSqlException exception) {
            Assert.assertEquals(FederationSqlErrorCode.QUERY_CANCELLED, exception.errorCode());
        } finally {
            registration.close();
            registry.close();
        }
    }

    /**
     * 验证 Core 登记前到达的取消不会被随后创建的 QueryState 覆盖。
     */
    @Test
    public void shouldRejectQueryBegunAfterPreCancellation() {
        QueryCancellationRegistry registry = new QueryCancellationRegistry();
        QueryId queryId = new QueryId("pre-cancelled-query");

        Assert.assertFalse(registry.cancel(queryId));
        try {
            registry.begin(queryId);
            Assert.fail("pre-cancelled query should not begin");
        } catch (FederationSqlException exception) {
            Assert.assertEquals(
                FederationSqlErrorCode.QUERY_CANCELLED,
                exception.errorCode()
            );
        } finally {
            registry.close();
        }
    }

    /**
     * 验证联邦查询登记的多个分片 Statement 会被同一次取消全部关闭。
     */
    @Test
    public void shouldCancelEveryRegisteredFragmentStatement() throws Exception {
        QueryCancellationRegistry registry = new QueryCancellationRegistry();
        QueryId queryId = new QueryId("federated-cancel");
        QueryCancellationRegistry.QueryRegistration registration = registry.begin(queryId);
        AtomicInteger firstCancel = new AtomicInteger();
        AtomicInteger firstClose = new AtomicInteger();
        AtomicInteger secondCancel = new AtomicInteger();
        AtomicInteger secondClose = new AtomicInteger();
        registration.statementLifecycle().register(statement(firstCancel, firstClose));
        registration.statementLifecycle().register(statement(secondCancel, secondClose));

        Assert.assertTrue(registry.cancel(queryId));
        awaitTermination(firstCancel, firstClose, 1);
        awaitTermination(secondCancel, secondClose, 1);

        registration.close();
        registry.close();
    }

    /**
     * 验证统一执行时限关闭全部 Statement，并保留 QUERY_TIMEOUT 终态。
     */
    @Test
    public void shouldRetainTimeoutReasonWhileTerminatingStatements() throws Exception {
        QueryCancellationRegistry registry = new QueryCancellationRegistry();
        QueryId queryId = new QueryId("timed-out-query");
        QueryCancellationRegistry.QueryRegistration registration = registry.begin(queryId);
        AtomicInteger cancelCalls = new AtomicInteger();
        AtomicInteger closeCalls = new AtomicInteger();
        StatementLifecycle lifecycle = registration.statementLifecycle();
        lifecycle.register(statement(cancelCalls, closeCalls));

        registration.requestTimeout();

        Assert.assertTrue(lifecycle.timeoutRequested());
        Assert.assertTrue(lifecycle.cancellationRequested());
        awaitTermination(cancelCalls, closeCalls, 1);
        try {
            registration.ensureNotCancelled();
            Assert.fail("timed out query must not continue");
        } catch (FederationSqlException exception) {
            Assert.assertEquals(FederationSqlErrorCode.QUERY_TIMEOUT, exception.errorCode());
        } finally {
            registration.close();
            registry.close();
        }
    }

    /**
     * 验证取消与超时并发到达时，查询只保留第一个成功写入的终止原因。
     *
     * @throws Exception 并发执行失败
     */
    @Test
    public void shouldKeepFirstTerminationReasonUnderCancelTimeoutRace() throws Exception {
        QueryCancellationRegistry registry = new QueryCancellationRegistry();
        QueryId queryId = new QueryId("termination-race");
        QueryCancellationRegistry.QueryRegistration registration = registry.begin(queryId);
        CountDownLatch start = new CountDownLatch(1);
        ExecutorService executor = Executors.newFixedThreadPool(2);
        try {
            var cancel = executor.submit(() -> {
                start.await();
                return registry.cancelOutcome(queryId).reason();
            });
            var timeout = executor.submit(() -> {
                start.await();
                return registration.requestTimeoutOutcome();
            });

            start.countDown();
            QueryCancellationRegistry.TerminationReason cancelObserved =
                cancel.get(2, TimeUnit.SECONDS);
            QueryCancellationRegistry.TerminationReason timeoutObserved =
                timeout.get(2, TimeUnit.SECONDS);

            Assert.assertEquals(cancelObserved, timeoutObserved);
            Assert.assertTrue(
                cancelObserved == QueryCancellationRegistry.TerminationReason.CANCELLED
                    || cancelObserved == QueryCancellationRegistry.TerminationReason.TIMED_OUT
            );
            try {
                registration.ensureNotCancelled();
                Assert.fail("terminated query must not continue");
            } catch (FederationSqlException exception) {
                FederationSqlErrorCode expected = cancelObserved
                    == QueryCancellationRegistry.TerminationReason.TIMED_OUT
                    ? FederationSqlErrorCode.QUERY_TIMEOUT
                    : FederationSqlErrorCode.QUERY_CANCELLED;
                Assert.assertEquals(expected, exception.errorCode());
            }
        } finally {
            registration.close();
            executor.shutdownNow();
            registry.close();
        }
    }

    /**
     * 验证终止队列饱和时转交隔离执行器，取消线程不会被 JDBC 清理反向阻塞。
     *
     * @throws Exception 并发测试失败
     */
    @Test
    public void shouldIsolateStatementTerminationWhenPrimaryQueueIsFull() throws Exception {
        ThreadPoolExecutor terminationExecutor = executor("cancel-overflow-test");
        ThreadPoolExecutor cleanupExecutor = executor("cleanup-overflow-test");
        ThreadPoolExecutor overflowExecutor = executor("resource-overflow-test");
        QueryCancellationRegistry registry = new QueryCancellationRegistry(
            terminationExecutor,
            cleanupExecutor,
            overflowExecutor
        );
        CountDownLatch workerStarted = new CountDownLatch(1);
        CountDownLatch releaseWorker = new CountDownLatch(1);
        terminationExecutor.execute(() -> {
            workerStarted.countDown();
            try {
                releaseWorker.await();
            } catch (InterruptedException exception) {
                Thread.currentThread().interrupt();
            }
        });
        Assert.assertTrue(workerStarted.await(2, TimeUnit.SECONDS));
        terminationExecutor.execute(() -> {
        });

        QueryId queryId = new QueryId("cancel-queue-overflow");
        QueryCancellationRegistry.QueryRegistration registration = registry.begin(queryId);
        AtomicInteger cancelCalls = new AtomicInteger();
        AtomicInteger closeCalls = new AtomicInteger();
        registration.statementLifecycle().register(statement(cancelCalls, closeCalls));

        long startedAt = System.nanoTime();
        Assert.assertTrue(registry.cancel(queryId));
        Assert.assertTrue(
            TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startedAt) < 200L
        );
        awaitTermination(cancelCalls, closeCalls, 1);
        Assert.assertEquals(1L, registry.overflowFallbackCount());
        Assert.assertEquals(0L, registry.unscheduledCleanupCount());

        releaseWorker.countDown();
        registration.close();
        registry.close();
    }

    /**
     * 验证游标清理队列饱和时转交隔离执行器，并明确返回主队列未接收。
     *
     * @throws Exception 并发测试失败
     */
    @Test
    public void shouldIsolateCursorCleanupWhenPrimaryQueueIsFull() throws Exception {
        ThreadPoolExecutor terminationExecutor = executor("cancel-cleanup-test");
        ThreadPoolExecutor cleanupExecutor = executor("cleanup-overflow-test");
        ThreadPoolExecutor overflowExecutor = executor("resource-overflow-test");
        QueryCancellationRegistry registry = new QueryCancellationRegistry(
            terminationExecutor,
            cleanupExecutor,
            overflowExecutor
        );
        CountDownLatch workerStarted = new CountDownLatch(1);
        CountDownLatch releaseWorker = new CountDownLatch(1);
        cleanupExecutor.execute(() -> {
            workerStarted.countDown();
            try {
                releaseWorker.await();
            } catch (InterruptedException exception) {
                Thread.currentThread().interrupt();
            }
        });
        Assert.assertTrue(workerStarted.await(2, TimeUnit.SECONDS));
        cleanupExecutor.execute(() -> {
        });
        AtomicInteger cleanupCalls = new AtomicInteger();

        Assert.assertFalse(registry.submitCleanup(cleanupCalls::incrementAndGet));
        awaitValue(cleanupCalls, 1);
        Assert.assertEquals(1L, registry.overflowFallbackCount());
        Assert.assertEquals(0L, registry.unscheduledCleanupCount());

        releaseWorker.countDown();
        registry.close();
    }

    /**
     * 验证主队列和隔离队列同时饱和时，清理进入有界延期队列并在容量恢复后重试。
     *
     * @throws Exception 并发测试失败
     */
    @Test
    public void shouldRetryDeferredCleanupAfterBothExecutorsRecover() throws Exception {
        ThreadPoolExecutor terminationExecutor = executor("cancel-deferred-test");
        ThreadPoolExecutor cleanupExecutor = executor("cleanup-deferred-test");
        ThreadPoolExecutor overflowExecutor = executor("overflow-deferred-test");
        QueryCancellationRegistry registry = new QueryCancellationRegistry(
            terminationExecutor,
            cleanupExecutor,
            overflowExecutor
        );
        CountDownLatch cleanupStarted = new CountDownLatch(1);
        CountDownLatch overflowStarted = new CountDownLatch(1);
        CountDownLatch releaseCleanup = new CountDownLatch(1);
        CountDownLatch releaseOverflow = new CountDownLatch(1);
        blockExecutor(cleanupExecutor, cleanupStarted, releaseCleanup);
        blockExecutor(overflowExecutor, overflowStarted, releaseOverflow);
        Assert.assertTrue(cleanupStarted.await(2, TimeUnit.SECONDS));
        Assert.assertTrue(overflowStarted.await(2, TimeUnit.SECONDS));
        cleanupExecutor.execute(() -> { });
        overflowExecutor.execute(() -> { });
        AtomicInteger cleanupCalls = new AtomicInteger();

        Assert.assertFalse(registry.submitCleanup(cleanupCalls::incrementAndGet));
        Assert.assertEquals(1L, registry.cleanupMetrics().deferredRetries());
        Assert.assertEquals(1, registry.cleanupMetrics().deferredQueueDepth());

        releaseOverflow.countDown();
        awaitValue(cleanupCalls, 1);
        Assert.assertEquals(0, registry.cleanupMetrics().deferredQueueDepth());
        Assert.assertEquals(0L, registry.cleanupMetrics().unresolvedCleanups());

        releaseCleanup.countDown();
        registry.close();
    }

    /**
     * 验证 Driver 清理忽略中断时，登记器关闭仍保持有界且不会卡死调用线程。
     *
     * @throws Exception 并发测试失败
     */
    @Test
    public void shouldCloseWithinBoundWhenDriverCleanupBlocks() throws Exception {
        ThreadPoolExecutor terminationExecutor = executor("cancel-close-test");
        ThreadPoolExecutor cleanupExecutor = executor("cleanup-close-test");
        ThreadPoolExecutor overflowExecutor = executor("overflow-close-test");
        QueryCancellationRegistry registry = new QueryCancellationRegistry(
            terminationExecutor,
            cleanupExecutor,
            overflowExecutor
        );
        CountDownLatch releaseDriver = new CountDownLatch(1);
        QueryCancellationRegistry.QueryRegistration first = registry.begin(
            new QueryId("blocked-driver-1")
        );
        QueryCancellationRegistry.QueryRegistration second = registry.begin(
            new QueryId("blocked-driver-2")
        );
        first.statementLifecycle().register(blockingStatement(releaseDriver));
        second.statementLifecycle().register(blockingStatement(releaseDriver));
        registry.cancel(new QueryId("blocked-driver-1"));
        registry.cancel(new QueryId("blocked-driver-2"));

        long startedAt = System.nanoTime();
        try {
            registry.close();
            long elapsedMillis = TimeUnit.NANOSECONDS.toMillis(
                System.nanoTime() - startedAt
            );
            Assert.assertTrue("close took " + elapsedMillis + "ms", elapsedMillis < 1_000L);
            Assert.assertTrue(registry.cleanupMetrics().unresolvedCleanups() > 0L);
        } finally {
            releaseDriver.countDown();
            first.close();
            second.close();
        }
    }

    private static void blockExecutor(
        ThreadPoolExecutor executor,
        CountDownLatch started,
        CountDownLatch release
    ) {
        executor.execute(() -> {
            started.countDown();
            try {
                release.await();
            } catch (InterruptedException exception) {
                Thread.currentThread().interrupt();
            }
        });
    }

    private static void awaitTermination(
        AtomicInteger cancelCalls,
        AtomicInteger closeCalls,
        int expected
    ) throws InterruptedException {
        long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2);
        while ((cancelCalls.get() != expected || closeCalls.get() != expected)
            && System.nanoTime() - deadline < 0L) {
            Thread.sleep(5L);
        }
        Assert.assertEquals(expected, cancelCalls.get());
        Assert.assertEquals(expected, closeCalls.get());
    }

    private static void awaitValue(AtomicInteger value, int expected)
        throws InterruptedException {
        long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2);
        while (value.get() != expected && System.nanoTime() - deadline < 0L) {
            Thread.sleep(5L);
        }
        Assert.assertEquals(expected, value.get());
    }

    private static Statement blockingStatement(CountDownLatch releaseDriver) {
        return (Statement) Proxy.newProxyInstance(
            QueryCancellationRegistryTest.class.getClassLoader(),
            new Class<?>[] {Statement.class},
            (proxy, method, arguments) -> {
                if ("cancel".equals(method.getName())) {
                    while (releaseDriver.getCount() > 0L) {
                        try {
                            releaseDriver.await();
                        } catch (InterruptedException ignored) {
                            // 模拟忽略线程中断的 JDBC Driver。
                        }
                    }
                }
                return defaultValue(method.getReturnType());
            }
        );
    }

    private static Statement statement(
        AtomicInteger cancelCalls,
        AtomicInteger closeCalls
    ) {
        return (Statement) Proxy.newProxyInstance(
            QueryCancellationRegistryTest.class.getClassLoader(),
            new Class<?>[] {Statement.class},
            (proxy, method, arguments) -> {
                if ("cancel".equals(method.getName())) {
                    cancelCalls.incrementAndGet();
                }
                if ("close".equals(method.getName())) {
                    closeCalls.incrementAndGet();
                }
                return defaultValue(method.getReturnType());
            }
        );
    }

    private static ThreadPoolExecutor executor(String threadName) {
        return new ThreadPoolExecutor(
            1,
            1,
            0L,
            TimeUnit.MILLISECONDS,
            new ArrayBlockingQueue<>(1),
            runnable -> {
                Thread thread = new Thread(runnable, threadName);
                thread.setDaemon(true);
                return thread;
            },
            new ThreadPoolExecutor.AbortPolicy()
        );
    }

    private static Object defaultValue(Class<?> type) {
        if (!type.isPrimitive()) {
            return null;
        }
        if (type == boolean.class) {
            return false;
        }
        if (type == byte.class) {
            return (byte) 0;
        }
        if (type == short.class) {
            return (short) 0;
        }
        if (type == int.class) {
            return 0;
        }
        if (type == long.class) {
            return 0L;
        }
        if (type == float.class) {
            return 0F;
        }
        if (type == double.class) {
            return 0D;
        }
        if (type == char.class) {
            return '\0';
        }
        return null;
    }
}
