package com.easyagents.federation.sql.runtime;

import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import java.time.Duration;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import org.junit.Assert;
import org.junit.Test;

/**
 * 节点联邦内存准入的并发与释放契约测试。
 */
public class NodeMemoryAdmissionControllerTest {

    /**
     * 验证单个查询预算超过节点上限时立即拒绝，不进入等待队列。
     */
    @Test
    public void shouldRejectRequestLargerThanNodeBudget() {
        NodeMemoryAdmissionController controller = new NodeMemoryAdmissionController(10L);

        FederationSqlException failure = expectFailure(() -> controller.acquire(
            11L,
            QueryDeadline.compile(Duration.ofSeconds(1))
        ));

        Assert.assertEquals(
            FederationSqlErrorCode.FEDERATION_RESOURCE_LIMIT_EXCEEDED,
            failure.errorCode()
        );
    }

    /**
     * 验证聚合预留达到节点上限后，等待者受同一个绝对截止时间约束。
     */
    @Test
    public void shouldRespectDeadlineWhileWaitingForAggregateBudget() {
        NodeMemoryAdmissionController controller = new NodeMemoryAdmissionController(10L);
        try (NodeMemoryAdmissionController.Permit ignored = controller.acquire(
            6L,
            QueryDeadline.compile(Duration.ofSeconds(1))
        )) {
            FederationSqlException failure = expectFailure(() -> controller.acquire(
                5L,
                QueryDeadline.compile(Duration.ofMillis(20))
            ));

            Assert.assertEquals(
                FederationSqlErrorCode.NODE_MEMORY_ADMISSION_TIMEOUT,
                failure.errorCode()
            );
        }
    }

    /**
     * 验证准入等待收到取消时保留 QUERY_CANCELLED，不误报资源等待超时。
     */
    @Test
    public void shouldPreserveCancellationDuringAdmissionWait() {
        NodeMemoryAdmissionController controller = new NodeMemoryAdmissionController(10L);
        QueryCancellationRegistry registry = new QueryCancellationRegistry();
        var registration = registry.begin(new com.easyagents.federation.sql.execute.QueryId(
            "memory-admission-cancel"
        ));
        try (NodeMemoryAdmissionController.Permit ignored = controller.acquire(
            10L,
            QueryDeadline.compile(Duration.ofSeconds(1))
        )) {
            registry.cancel(new com.easyagents.federation.sql.execute.QueryId(
                "memory-admission-cancel"
            ));
            FederationSqlException failure = expectFailure(() -> controller.acquire(
                1L,
                QueryDeadline.query(
                    QueryDeadline.deadlineAfter(Duration.ofSeconds(1)),
                    registration
                )
            ));

            Assert.assertEquals(FederationSqlErrorCode.QUERY_CANCELLED, failure.errorCode());
        } finally {
            registration.close();
            registry.close();
        }
    }

    /**
     * 验证重复关闭不会增加许可，等待者只在真实预算释放后获得准入。
     *
     * @throws Exception 并发测试失败
     */
    @Test
    public void shouldReleaseMemoryPermitExactlyOnce() throws Exception {
        NodeMemoryAdmissionController controller = new NodeMemoryAdmissionController(10L);
        ExecutorService executor = Executors.newSingleThreadExecutor();
        try {
            NodeMemoryAdmissionController.Permit first = controller.acquire(
                10L,
                QueryDeadline.compile(Duration.ofSeconds(2))
            );
            first.close();
            first.close();

            NodeMemoryAdmissionController.Permit second = controller.acquire(
                10L,
                QueryDeadline.compile(Duration.ofSeconds(2))
            );
            Future<NodeMemoryAdmissionController.Permit> waiting = executor.submit(() ->
                controller.acquire(10L, QueryDeadline.compile(Duration.ofSeconds(2)))
            );

            TimeUnit.MILLISECONDS.sleep(75L);
            Assert.assertFalse(waiting.isDone());
            second.close();

            NodeMemoryAdmissionController.Permit third = waiting.get(1, TimeUnit.SECONDS);
            third.close();
        } finally {
            executor.shutdownNow();
        }
    }

    /**
     * 验证线程中断与资源等待超时采用不同错误码，并保留线程中断标记。
     *
     * @throws Exception 并发测试失败
     */
    @Test
    public void shouldReportExecutionFailureWhenAdmissionWaitIsInterrupted() throws Exception {
        NodeMemoryAdmissionController controller = new NodeMemoryAdmissionController(10L);
        AtomicReference<FederationSqlException> failure = new AtomicReference<>();
        AtomicReference<Boolean> interrupted = new AtomicReference<>(false);
        Thread waiter;
        try (NodeMemoryAdmissionController.Permit ignored = controller.acquire(
            10L,
            QueryDeadline.compile(Duration.ofSeconds(2))
        )) {
            waiter = new Thread(() -> {
                try {
                    controller.acquire(1L, QueryDeadline.compile(Duration.ofSeconds(2)));
                } catch (FederationSqlException exception) {
                    failure.set(exception);
                    interrupted.set(Thread.currentThread().isInterrupted());
                }
            });
            waiter.start();
            TimeUnit.MILLISECONDS.sleep(50L);
            waiter.interrupt();
            waiter.join(1_000L);
        }

        Assert.assertFalse(waiter.isAlive());
        Assert.assertNotNull(failure.get());
        Assert.assertEquals(
            FederationSqlErrorCode.EXECUTION_FAILED,
            failure.get().errorCode()
        );
        Assert.assertTrue(interrupted.get());
    }

    /**
     * 执行预期失败的内存准入动作。
     *
     * @param action 测试动作
     * @return 捕获的统一异常
     */
    private static FederationSqlException expectFailure(Runnable action) {
        try {
            action.run();
            Assert.fail("expected memory admission to fail");
            return null;
        } catch (FederationSqlException exception) {
            return exception;
        }
    }
}
