package com.easyagents.federation.sql.runtime;

import com.easyagents.federation.sql.adapter.AdapterCompatibility;
import com.easyagents.federation.sql.adapter.AdapterCompatibilityStatus;
import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import com.easyagents.federation.sql.compile.FederationSqlPlan;
import com.easyagents.federation.sql.source.SourceId;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
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.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.time.Duration;
import java.time.Instant;
import org.junit.Assert;
import org.junit.Test;

/**
 * 有界计划缓存容量与 single-flight 测试。
 */
public class BoundedPlanCacheTest {

    /**
     * 验证唯一 SQL 不会让缓存无界增长。
     */
    @Test
    public void shouldRemainBounded() {
        BoundedPlanCache cache = new BoundedPlanCache(4);
        for (int index = 0; index < 20; index++) {
            PlanCacheKey key = key("select " + index);
            cache.getOrCompile(key, BoundedPlanCacheTest::plan);
        }
        Assert.assertEquals(4, cache.size());
        cache.close();
    }

    /**
     * 验证相同冷键的并发请求只执行一次编译。
     *
     * @throws Exception 并发测试失败
     */
    @Test
    public void shouldCompileSameColdKeyOnce() throws Exception {
        BoundedPlanCache cache = new BoundedPlanCache(4);
        PlanCacheKey key = key("select 1");
        AtomicInteger compilations = new AtomicInteger();
        int concurrency = 20;
        ExecutorService executor = Executors.newFixedThreadPool(8);
        CountDownLatch start = new CountDownLatch(1);
        CountDownLatch done = new CountDownLatch(concurrency);
        for (int index = 0; index < concurrency; index++) {
            executor.execute(() -> {
                try {
                    start.await();
                    cache.getOrCompile(key, () -> {
                        compilations.incrementAndGet();
                        return plan();
                    });
                } catch (InterruptedException exception) {
                    Thread.currentThread().interrupt();
                } finally {
                    done.countDown();
                }
            });
        }
        start.countDown();
        Assert.assertTrue(done.await(5, TimeUnit.SECONDS));
        executor.shutdownNow();
        Assert.assertEquals(1, compilations.get());
        cache.close();
    }

    /**
     * 验证关闭期间的冷编译不会在 close 返回后重新写入缓存。
     *
     * @throws Exception 并发测试失败
     */
    @Test
    public void shouldRejectCompilationCompletedAfterClose() throws Exception {
        BoundedPlanCache cache = new BoundedPlanCache(2);
        CountDownLatch compiling = new CountDownLatch(1);
        CountDownLatch release = new CountDownLatch(1);
        ExecutorService executor = Executors.newSingleThreadExecutor();
        try {
            Future<FederationSqlErrorCode> future = executor.submit(() -> {
                try {
                    cache.getOrCompile(key("select slow"), () -> {
                        compiling.countDown();
                        try {
                            release.await();
                        } catch (InterruptedException exception) {
                            Thread.currentThread().interrupt();
                            throw new IllegalStateException(exception);
                        }
                        return plan();
                    });
                    return null;
                } catch (FederationSqlException exception) {
                    return exception.errorCode();
                }
            });
            Assert.assertTrue(compiling.await(2, TimeUnit.SECONDS));
            cache.close();
            release.countDown();
            Assert.assertEquals(FederationSqlErrorCode.ENGINE_CLOSED, future.get(2, TimeUnit.SECONDS));
            Assert.assertEquals(0, cache.size());
        } finally {
            release.countDown();
            executor.shutdownNow();
            cache.close();
        }
    }

    /**
     * 验证冷编译并发上限与缓存容量相互独立。
     *
     * @throws Exception 并发测试失败
     */
    @Test
    public void shouldLimitConcurrentColdCompilations() throws Exception {
        BoundedPlanCache cache = new BoundedPlanCache(16, 2);
        AtomicInteger active = new AtomicInteger();
        AtomicInteger maximumActive = new AtomicInteger();
        CountDownLatch twoCompilersEntered = new CountDownLatch(2);
        CountDownLatch release = new CountDownLatch(1);
        ExecutorService executor = Executors.newFixedThreadPool(6);
        List<Future<FederationSqlPlan>> futures = new ArrayList<>();
        try {
            for (int index = 0; index < 6; index++) {
                int sqlIndex = index;
                futures.add(executor.submit(() -> cache.getOrCompile(key("select " + sqlIndex), () -> {
                    int current = active.incrementAndGet();
                    maximumActive.accumulateAndGet(current, Math::max);
                    twoCompilersEntered.countDown();
                    try {
                        release.await();
                        return plan();
                    } catch (InterruptedException exception) {
                        Thread.currentThread().interrupt();
                        throw new IllegalStateException(exception);
                    } finally {
                        active.decrementAndGet();
                    }
                })));
            }
            Assert.assertTrue(twoCompilersEntered.await(2, TimeUnit.SECONDS));
            Assert.assertEquals(2, maximumActive.get());
            release.countDown();
            for (Future<FederationSqlPlan> future : futures) {
                Assert.assertNotNull(future.get(2, TimeUnit.SECONDS));
            }
            Assert.assertEquals(2, maximumActive.get());
        } finally {
            release.countDown();
            executor.shutdownNow();
            cache.close();
        }
    }

    /**
     * 验证超出权重的计划不进入缓存，条目权重始终保持有界。
     */
    @Test
    public void shouldRejectOversizedPlanFromCache() {
        BoundedPlanCache cache = new BoundedPlanCache(4, 2, 1, Duration.ofMinutes(1));
        AtomicInteger compilations = new AtomicInteger();
        PlanCacheKey key = key("select weighted");

        cache.getOrCompile(key, () -> {
            compilations.incrementAndGet();
            return plan();
        });
        cache.getOrCompile(key, () -> {
            compilations.incrementAndGet();
            return plan();
        });

        Assert.assertEquals(2, compilations.get());
        Assert.assertEquals(0, cache.size());
        Assert.assertEquals(0, cache.weightBytes());
        cache.close();
    }

    /**
     * 验证 TTL 到期后旧计划不会继续命中。
     *
     * @throws Exception 休眠被中断
     */
    @Test
    public void shouldRecompileExpiredPlan() throws Exception {
        BoundedPlanCache cache = new BoundedPlanCache(
            4,
            2,
            1_000_000,
            Duration.ofMillis(10)
        );
        AtomicInteger compilations = new AtomicInteger();
        PlanCacheKey key = key("select expiring");
        cache.getOrCompile(key, () -> {
            compilations.incrementAndGet();
            return plan();
        });
        Thread.sleep(30);
        cache.getOrCompile(key, () -> {
            compilations.incrementAndGet();
            return plan();
        });

        Assert.assertEquals(2, compilations.get());
        Assert.assertEquals(1, cache.size());
        cache.close();
    }

    /**
     * 验证单调时钟跨越 long 回绕边界时，TTL 仍按经过时长判断。
     */
    @Test
    public void shouldExpireByElapsedTimeAcrossNanoTimeWrap() {
        AtomicLong clock = new AtomicLong(Long.MAX_VALUE - 5L);
        BoundedPlanCache cache = new BoundedPlanCache(
            4,
            2,
            1_000_000,
            Duration.ofNanos(10L),
            clock::get
        );
        AtomicInteger compilations = new AtomicInteger();
        PlanCacheKey key = key("select wrap");

        cache.getOrCompile(key, () -> {
            compilations.incrementAndGet();
            return plan();
        });
        clock.set(Long.MIN_VALUE + 1L);
        cache.getOrCompile(key, () -> {
            compilations.incrementAndGet();
            return plan();
        });
        Assert.assertEquals(1, compilations.get());

        clock.set(Long.MIN_VALUE + 6L);
        cache.getOrCompile(key, () -> {
            compilations.incrementAndGet();
            return plan();
        });
        Assert.assertEquals(2, compilations.get());
        cache.close();
    }

    /**
     * 验证统计快照先于缓存 TTL 失效时旧成本计划会重新编译。
     *
     * @throws Exception 休眠被中断
     */
    @Test
    public void shouldNotCacheBeyondStatisticsValidity() throws Exception {
        BoundedPlanCache cache = new BoundedPlanCache(
            4,
            2,
            1_000_000,
            Duration.ofMinutes(1)
        );
        AtomicInteger compilations = new AtomicInteger();
        PlanCacheKey key = key("select statistics expiring");
        cache.getOrCompile(key, () -> {
            compilations.incrementAndGet();
            return plan(Instant.now().plusMillis(10));
        });
        Thread.sleep(30);
        cache.getOrCompile(key, () -> {
            compilations.incrementAndGet();
            return plan(Instant.MAX);
        });

        Assert.assertEquals(2, compilations.get());
        Assert.assertEquals(1, cache.size());
        cache.close();
    }

    /**
     * 验证关闭 Runtime 时只移除引用其完整身份的计划。
     */
    @Test
    public void shouldInvalidateExactRuntimeIdentity() {
        BoundedPlanCache cache = new BoundedPlanCache(4);
        cache.getOrCompile(key("select exact"), BoundedPlanCacheTest::plan);
        Assert.assertEquals(1, cache.size());

        cache.invalidateRuntimeIdentity(
            new SourceId("source"),
            1,
            "checksum",
            "fingerprint"
        );

        Assert.assertEquals(0, cache.size());
        cache.close();
    }

    /**
     * 验证调用方判定查询级统计已变化时，缓存会丢弃旧计划并重新编译。
     */
    @Test
    public void shouldRecompileWhenContextValidatorRejectsCachedPlan() {
        BoundedPlanCache cache = new BoundedPlanCache(4);
        PlanCacheKey key = key("select context scoped statistics");
        AtomicInteger compilations = new AtomicInteger();
        FederationSqlPlan first = cache.getOrCompileWithStatus(
            key,
            () -> {
                compilations.incrementAndGet();
                return plan();
            },
            BoundedPlanCache.WaitGuard.unbounded(),
            ignored -> true
        ).plan();

        BoundedPlanCache.LookupResult refreshed = cache.getOrCompileWithStatus(
            key,
            () -> {
                compilations.incrementAndGet();
                return plan();
            },
            BoundedPlanCache.WaitGuard.unbounded(),
            candidate -> candidate != first
        );

        Assert.assertEquals(2, compilations.get());
        Assert.assertFalse(refreshed.cacheHit());
        Assert.assertNotSame(first, refreshed.plan());
        cache.close();
    }

    /**
     * 验证迟到的旧统计请求不会淘汰或覆盖已经缓存的新代计划。
     *
     * @throws Exception 并发测试失败
     */
    @Test
    public void shouldKeepNewerPlanWhenOlderGenerationFinishesLate() throws Exception {
        BoundedPlanCache cache = new BoundedPlanCache(4, 2);
        PlanCacheKey key = key("select generation protected");
        Instant olderGeneration = Instant.parse("2026-08-23T00:00:00Z");
        Instant newerGeneration = olderGeneration.plusSeconds(1);
        FederationSqlPlan olderPlan = plan();
        FederationSqlPlan newerPlan = plan();
        cache.getOrCompileWithStatus(
            key,
            () -> newerPlan,
            BoundedPlanCache.WaitGuard.unbounded(),
            candidate -> candidate == newerPlan,
            newerGeneration
        );
        CountDownLatch olderCompiling = new CountDownLatch(1);
        CountDownLatch releaseOlder = new CountDownLatch(1);
        ExecutorService executor = Executors.newSingleThreadExecutor();
        try {
            Future<BoundedPlanCache.LookupResult> olderLookup = executor.submit(() ->
                cache.getOrCompileWithStatus(
                    key,
                    () -> {
                        olderCompiling.countDown();
                        awaitLatch(releaseOlder);
                        return olderPlan;
                    },
                    BoundedPlanCache.WaitGuard.unbounded(),
                    candidate -> candidate == olderPlan,
                    olderGeneration
                )
            );
            Assert.assertTrue(olderCompiling.await(2, TimeUnit.SECONDS));

            BoundedPlanCache.LookupResult concurrentNewerHit = cache.getOrCompileWithStatus(
                key,
                () -> {
                    throw new AssertionError("newer plan should remain cached");
                },
                BoundedPlanCache.WaitGuard.unbounded(),
                candidate -> candidate == newerPlan,
                newerGeneration
            );
            Assert.assertTrue(concurrentNewerHit.cacheHit());
            Assert.assertSame(newerPlan, concurrentNewerHit.plan());

            releaseOlder.countDown();
            Assert.assertSame(olderPlan, olderLookup.get(2, TimeUnit.SECONDS).plan());
            BoundedPlanCache.LookupResult finalNewerHit = cache.getOrCompileWithStatus(
                key,
                () -> {
                    throw new AssertionError("older plan must not replace newer generation");
                },
                BoundedPlanCache.WaitGuard.unbounded(),
                candidate -> candidate == newerPlan,
                newerGeneration
            );
            Assert.assertTrue(finalNewerHit.cacheHit());
            Assert.assertSame(newerPlan, finalNewerHit.plan());
            Assert.assertEquals(1, cache.size());
        } finally {
            releaseOlder.countDown();
            executor.shutdownNow();
            cache.close();
        }
    }

    /**
     * 验证编译完成但尚未发布时关闭缓存，迟到任务不能发布成功计划。
     *
     * @throws Exception 并发测试失败
     */
    @Test
    public void shouldRejectLateInFlightPublicationAfterClose() throws Exception {
        BoundedPlanCache cache = new BoundedPlanCache(2, 1);
        CountDownLatch readyToPublish = new CountDownLatch(1);
        CountDownLatch releasePublication = new CountDownLatch(1);
        ExecutorService executor = Executors.newSingleThreadExecutor();
        try {
            Future<FederationSqlErrorCode> publisher = executor.submit(() -> {
                try {
                    cache.getOrCompileWithStatus(
                        key("select late publication"),
                        BoundedPlanCacheTest::plan,
                        BoundedPlanCache.WaitGuard.unbounded(),
                        candidate -> {
                            readyToPublish.countDown();
                            awaitLatch(releasePublication);
                            return true;
                        },
                        Instant.parse("2026-08-23T00:00:00Z")
                    );
                    return null;
                } catch (FederationSqlException exception) {
                    return exception.errorCode();
                }
            });
            Assert.assertTrue(readyToPublish.await(2, TimeUnit.SECONDS));

            cache.close();
            releasePublication.countDown();

            Assert.assertEquals(
                FederationSqlErrorCode.ENGINE_CLOSED,
                publisher.get(2, TimeUnit.SECONDS)
            );
            Assert.assertEquals(0, cache.size());
            Assert.assertEquals(1, cache.availableCompileSlots());
        } finally {
            releasePublication.countDown();
            executor.shutdownNow();
            cache.close();
        }
    }

    /**
     * 验证可能读取宽表统计的上下文校验不会占用缓存全局锁。
     *
     * @throws Exception 并发测试失败
     */
    @Test
    public void shouldValidateCachedPlanOutsideGlobalLock() throws Exception {
        BoundedPlanCache cache = new BoundedPlanCache(4);
        PlanCacheKey key = key("select validation outside lock");
        cache.getOrCompile(key, BoundedPlanCacheTest::plan);
        CountDownLatch validating = new CountDownLatch(1);
        CountDownLatch release = new CountDownLatch(1);
        ExecutorService executor = Executors.newFixedThreadPool(2);
        try {
            Future<BoundedPlanCache.LookupResult> lookup = executor.submit(() ->
                cache.getOrCompileWithStatus(
                    key,
                    BoundedPlanCacheTest::plan,
                    BoundedPlanCache.WaitGuard.unbounded(),
                    candidate -> {
                        validating.countDown();
                        try {
                            release.await();
                        } catch (InterruptedException exception) {
                            Thread.currentThread().interrupt();
                            throw new IllegalStateException(exception);
                        }
                        return true;
                    }
                )
            );
            Assert.assertTrue(validating.await(2, TimeUnit.SECONDS));

            Future<Integer> size = executor.submit(cache::size);
            Assert.assertEquals(Integer.valueOf(1), size.get(500, TimeUnit.MILLISECONDS));

            release.countDown();
            Assert.assertTrue(lookup.get(2, TimeUnit.SECONDS).cacheHit());
        } finally {
            release.countDown();
            executor.shutdownNow();
            cache.close();
        }
    }

    /**
     * 验证同键 single-flight 等待会响应统一编译截止时间。
     *
     * @throws Exception 并发测试失败
     */
    @Test
    public void shouldStopWaitingForSameKeyWhenCompileDeadlineExpires() throws Exception {
        BoundedPlanCache cache = new BoundedPlanCache(4, 1);
        CountDownLatch compiling = new CountDownLatch(1);
        CountDownLatch release = new CountDownLatch(1);
        ExecutorService executor = Executors.newSingleThreadExecutor();
        try {
            Future<FederationSqlPlan> compiler = executor.submit(() -> cache.getOrCompile(
                key("select guarded"),
                () -> {
                    compiling.countDown();
                    try {
                        release.await();
                    } catch (InterruptedException exception) {
                        Thread.currentThread().interrupt();
                        throw new IllegalStateException(exception);
                    }
                    return plan();
                }
            ));
            Assert.assertTrue(compiling.await(2, TimeUnit.SECONDS));

            long deadline = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(30);
            try {
                cache.getOrCompileWithStatus(
                    key("select guarded"),
                    BoundedPlanCacheTest::plan,
                    deadlineGuard(deadline)
                );
                Assert.fail("single-flight waiter should respect compile deadline");
            } catch (FederationSqlException exception) {
                Assert.assertEquals(
                    FederationSqlErrorCode.SQL_COMPILE_TIMEOUT,
                    exception.errorCode()
                );
            }

            release.countDown();
            Assert.assertNotNull(compiler.get(2, TimeUnit.SECONDS));
        } finally {
            release.countDown();
            executor.shutdownNow();
            cache.close();
        }
    }

    /**
     * 验证冷编译槽等待同样受查询截止时间约束。
     *
     * @throws Exception 并发测试失败
     */
    @Test
    public void shouldStopWaitingForCompileSlotWhenDeadlineExpires() throws Exception {
        BoundedPlanCache cache = new BoundedPlanCache(4, 1);
        CountDownLatch compiling = new CountDownLatch(1);
        CountDownLatch release = new CountDownLatch(1);
        ExecutorService executor = Executors.newSingleThreadExecutor();
        try {
            Future<FederationSqlPlan> compiler = executor.submit(() -> cache.getOrCompile(
                key("select first"),
                () -> {
                    compiling.countDown();
                    try {
                        release.await();
                    } catch (InterruptedException exception) {
                        Thread.currentThread().interrupt();
                        throw new IllegalStateException(exception);
                    }
                    return plan();
                }
            ));
            Assert.assertTrue(compiling.await(2, TimeUnit.SECONDS));

            long deadline = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(30);
            try {
                cache.getOrCompileWithStatus(
                    key("select second"),
                    BoundedPlanCacheTest::plan,
                    deadlineGuard(deadline)
                );
                Assert.fail("compile-slot waiter should respect compile deadline");
            } catch (FederationSqlException exception) {
                Assert.assertEquals(
                    FederationSqlErrorCode.SQL_COMPILE_TIMEOUT,
                    exception.errorCode()
                );
            }

            release.countDown();
            Assert.assertNotNull(compiler.get(2, TimeUnit.SECONDS));
        } finally {
            release.countDown();
            executor.shutdownNow();
            cache.close();
        }
    }

    /**
     * 验证等待冷编译槽的线程遇到 close 后退出，随后归还的许可不会泄漏。
     *
     * @throws Exception 并发测试失败
     */
    @Test
    public void shouldReleaseCompileSlotWhenCacheClosesWhileAnotherThreadWaits() throws Exception {
        BoundedPlanCache cache = new BoundedPlanCache(4, 1);
        CountDownLatch compiling = new CountDownLatch(1);
        CountDownLatch releaseCompiler = new CountDownLatch(1);
        CountDownLatch waiterStarted = new CountDownLatch(1);
        ExecutorService executor = Executors.newFixedThreadPool(2);
        try {
            Future<FederationSqlErrorCode> compiler = executor.submit(() -> {
                try {
                    cache.getOrCompile(key("select holding slot"), () -> {
                        compiling.countDown();
                        try {
                            releaseCompiler.await();
                        } catch (InterruptedException exception) {
                            Thread.currentThread().interrupt();
                            throw new IllegalStateException(exception);
                        }
                        return plan();
                    });
                    return null;
                } catch (FederationSqlException exception) {
                    return exception.errorCode();
                }
            });
            Assert.assertTrue(compiling.await(2, TimeUnit.SECONDS));
            Future<FederationSqlErrorCode> waiter = executor.submit(() -> {
                waiterStarted.countDown();
                try {
                    cache.getOrCompile(key("select waiting for slot"), BoundedPlanCacheTest::plan);
                    return null;
                } catch (FederationSqlException exception) {
                    return exception.errorCode();
                }
            });
            Assert.assertTrue(waiterStarted.await(2, TimeUnit.SECONDS));

            cache.close();
            releaseCompiler.countDown();

            Assert.assertEquals(FederationSqlErrorCode.ENGINE_CLOSED, waiter.get(2, TimeUnit.SECONDS));
            Assert.assertEquals(FederationSqlErrorCode.ENGINE_CLOSED, compiler.get(2, TimeUnit.SECONDS));
            Assert.assertEquals(1, cache.availableCompileSlots());
        } finally {
            releaseCompiler.countDown();
            executor.shutdownNow();
            cache.close();
        }
    }

    private static BoundedPlanCache.WaitGuard deadlineGuard(long deadlineNanos) {
        return new BoundedPlanCache.WaitGuard() {
            @Override
            public void ensureAllowed() {
                if (remainingNanos() <= 0L) {
                    throw new FederationSqlException(
                        FederationSqlErrorCode.SQL_COMPILE_TIMEOUT,
                        "test compile deadline expired"
                    );
                }
            }

            @Override
            public long remainingNanos() {
                return Math.max(0L, deadlineNanos - System.nanoTime());
            }
        };
    }

    /**
     * 等待测试闩锁，并把中断转换为测试线程可见的失败。
     *
     * @param latch 测试闩锁
     */
    private static void awaitLatch(CountDownLatch latch) {
        try {
            latch.await();
        } catch (InterruptedException exception) {
            Thread.currentThread().interrupt();
            throw new IllegalStateException(exception);
        }
    }

    private static PlanCacheKey key(String sql) {
        return new PlanCacheKey(
            sql,
            new SourceId("source"),
            1,
            1,
            List.of(),
            "fake",
            "fingerprint",
            "engine-policy",
            "default",
            1
        );
    }

    private static FederationSqlPlan plan() {
        return plan(Instant.MAX);
    }

    private static FederationSqlPlan plan(Instant statisticsValidUntil) {
        return new DefaultFederationSqlPlan(
            new SourceId("source"),
            1,
            "SELECT 1",
            "SELECT 1",
            null,
            null,
            0,
            List.of(),
            List.of(),
            Set.of(new SourceId("source")),
            new AdapterCompatibility(
                AdapterCompatibilityStatus.VERIFIED,
                "fake",
                "1",
                "fake-driver",
                "1",
                "test"
            ),
            true,
            "checksum",
            "fake",
            "fingerprint",
            statisticsValidUntil
        );
    }
}
