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.compile.FederationSqlPlan;
import java.time.Duration;
import java.time.Instant;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.LockSupport;
import java.util.function.Predicate;
import java.util.function.LongSupplier;
import java.util.function.Supplier;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.RelVisitor;

/**
 * 有界 LRU 计划缓存，并对相同冷键执行 single-flight 编译。
 */
final class BoundedPlanCache implements AutoCloseable {

    private final int maximumEntries;
    private final int maximumConcurrentCompilations;
    private final long maximumWeightBytes;
    private final long ttlNanos;
    private final LongSupplier nanoTime;
    private final Map<PlanCacheKey, CacheEntry> entries;
    private final ConcurrentHashMap<PlanCacheKey, CompletableFuture<FederationSqlPlan>> inFlight =
        new ConcurrentHashMap<>();
    private final Semaphore compileSlots;
    private final AtomicBoolean closed = new AtomicBoolean();
    private long currentWeightBytes;

    /**
     * 创建有界计划缓存。
     *
     * @param maximumEntries 最大条目数
     */
    BoundedPlanCache(int maximumEntries) {
        this(
            maximumEntries,
            Math.min(maximumEntries, defaultCompilationConcurrency()),
            defaultMaximumWeight(maximumEntries),
            Duration.ofMinutes(30),
            System::nanoTime
        );
    }

    /**
     * 创建同时限制缓存容量与冷编译并发的计划缓存。
     *
     * @param maximumEntries 最大条目数
     * @param maximumConcurrentCompilations 最大并发冷编译数
     */
    BoundedPlanCache(int maximumEntries, int maximumConcurrentCompilations) {
        this(
            maximumEntries,
            maximumConcurrentCompilations,
            defaultMaximumWeight(maximumEntries),
            Duration.ofMinutes(30),
            System::nanoTime
        );
    }

    /**
     * 创建同时限制条目、估算权重、存活时间和冷编译并发的计划缓存。
     *
     * @param maximumEntries 最大条目数
     * @param maximumConcurrentCompilations 最大并发冷编译数
     * @param maximumWeightBytes 最大估算权重
     * @param timeToLive 条目存活时间
     */
    BoundedPlanCache(
        int maximumEntries,
        int maximumConcurrentCompilations,
        long maximumWeightBytes,
        Duration timeToLive
    ) {
        this(
            maximumEntries,
            maximumConcurrentCompilations,
            maximumWeightBytes,
            timeToLive,
            System::nanoTime
        );
    }

    /**
     * 创建使用指定单调时钟的计划缓存，供回绕边界测试使用。
     *
     * @param maximumEntries 最大条目数
     * @param maximumConcurrentCompilations 最大并发冷编译数
     * @param maximumWeightBytes 最大估算权重
     * @param timeToLive 条目存活时间
     * @param nanoTime 单调时钟
     */
    BoundedPlanCache(
        int maximumEntries,
        int maximumConcurrentCompilations,
        long maximumWeightBytes,
        Duration timeToLive,
        LongSupplier nanoTime
    ) {
        if (maximumEntries <= 0) {
            throw new IllegalArgumentException("maximumEntries must be positive");
        }
        if (maximumConcurrentCompilations <= 0) {
            throw new IllegalArgumentException("maximumConcurrentCompilations must be positive");
        }
        if (maximumWeightBytes <= 0) {
            throw new IllegalArgumentException("maximumWeightBytes must be positive");
        }
        if (timeToLive == null || timeToLive.isZero() || timeToLive.isNegative()) {
            throw new IllegalArgumentException("timeToLive must be positive");
        }
        this.maximumEntries = maximumEntries;
        this.maximumConcurrentCompilations = maximumConcurrentCompilations;
        this.maximumWeightBytes = maximumWeightBytes;
        this.ttlNanos = saturatingNanos(timeToLive);
        this.nanoTime = java.util.Objects.requireNonNull(nanoTime, "nanoTime");
        this.compileSlots = new Semaphore(maximumConcurrentCompilations, true);
        this.entries = new LinkedHashMap<>(16, 0.75f, true);
    }

    /**
     * 命中计划或由单个线程完成冷编译。
     *
     * @param key 缓存键
     * @param compiler 冷编译函数
     * @return 编译计划
     */
    FederationSqlPlan getOrCompile(PlanCacheKey key, Supplier<FederationSqlPlan> compiler) {
        return getOrCompileWithStatus(key, compiler, WaitGuard.unbounded()).plan();
    }

    /**
     * 命中计划或完成冷编译，并返回是否复用了缓存或并发 single-flight。
     *
     * @param key 缓存键
     * @param compiler 冷编译函数
     * @return 缓存查询结果
     */
    LookupResult getOrCompileWithStatus(
        PlanCacheKey key,
        Supplier<FederationSqlPlan> compiler
    ) {
        return getOrCompileWithStatus(key, compiler, WaitGuard.unbounded());
    }

    /**
     * 命中计划或完成冷编译，并让 single-flight 与编译槽等待响应查询终态。
     *
     * @param key 缓存键
     * @param compiler 冷编译函数
     * @param waitGuard 等待终态检查器
     * @return 缓存查询结果
     */
    LookupResult getOrCompileWithStatus(
        PlanCacheKey key,
        Supplier<FederationSqlPlan> compiler,
        WaitGuard waitGuard
    ) {
        return getOrCompileWithStatus(key, compiler, waitGuard, plan -> true);
    }

    /**
     * 命中计划或完成冷编译，并按调用方当前上下文校验缓存计划。
     *
     * @param key 缓存键
     * @param compiler 冷编译函数
     * @param waitGuard 等待终态检查器
     * @param cachedPlanValidator 缓存计划是否仍适用于当前上下文
     * @return 缓存查询结果
     */
    LookupResult getOrCompileWithStatus(
        PlanCacheKey key,
        Supplier<FederationSqlPlan> compiler,
        WaitGuard waitGuard,
        Predicate<FederationSqlPlan> cachedPlanValidator
    ) {
        return getOrCompileWithStatus(
            key,
            compiler,
            waitGuard,
            cachedPlanValidator,
            Instant.MIN
        );
    }

    /**
     * 命中计划或完成冷编译，并阻止较旧上下文淘汰或覆盖较新的缓存计划。
     *
     * @param key 缓存键
     * @param compiler 冷编译函数
     * @param waitGuard 等待终态检查器
     * @param cachedPlanValidator 缓存计划是否仍适用于当前上下文
     * @param contextGeneration 当前上下文的可比较代际
     * @return 缓存查询结果
     */
    LookupResult getOrCompileWithStatus(
        PlanCacheKey key,
        Supplier<FederationSqlPlan> compiler,
        WaitGuard waitGuard,
        Predicate<FederationSqlPlan> cachedPlanValidator,
        Instant contextGeneration
    ) {
        Instant generation = java.util.Objects.requireNonNull(
            contextGeneration,
            "contextGeneration"
        );
        while (true) {
            ensureOpen();
            waitGuard.ensureAllowed();
            CacheEntry cached;
            synchronized (entries) {
                cached = cachedEntry(key, nanoTime.getAsLong());
            }
            if (cached != null && cachedPlanValidator.test(cached.plan())) {
                return new LookupResult(cached.plan(), true);
            }
            if (cached != null) {
                synchronized (entries) {
                    removeEntryIfNotNewer(key, cached.plan(), generation);
                }
            }

            CompletableFuture<FederationSqlPlan> existing = inFlight.get(key);
            if (existing != null) {
                FederationSqlPlan sharedPlan = await(existing, waitGuard);
                if (cachedPlanValidator.test(sharedPlan)) {
                    return new LookupResult(sharedPlan, true);
                }
                synchronized (entries) {
                    removeEntryIfNotNewer(key, sharedPlan, generation);
                }
                awaitInFlightRemoval(key, existing, waitGuard);
                continue;
            }

            acquireCompileSlot(waitGuard);
            CompletableFuture<FederationSqlPlan> ownFuture = new CompletableFuture<>();
            try {
                ensureOpen();
                CacheEntry cachedAfterPermit;
                synchronized (entries) {
                    cachedAfterPermit = cachedEntry(key, nanoTime.getAsLong());
                }
                if (cachedAfterPermit != null
                    && cachedPlanValidator.test(cachedAfterPermit.plan())) {
                    compileSlots.release();
                    return new LookupResult(cachedAfterPermit.plan(), true);
                }
                if (cachedAfterPermit != null) {
                    synchronized (entries) {
                        removeEntryIfNotNewer(
                            key,
                            cachedAfterPermit.plan(),
                            generation
                        );
                    }
                }
                synchronized (entries) {
                    // 与 close 使用同一线性化锁，禁止关闭后登记新的编译任务。
                    ensureOpen();
                    existing = inFlight.putIfAbsent(key, ownFuture);
                }
            } catch (Throwable throwable) {
                compileSlots.release();
                throw throwable;
            }
            if (existing != null) {
                // 相同键只短暂占用配额完成握手，等待期间不阻塞其他冷键。
                compileSlots.release();
                FederationSqlPlan sharedPlan = await(existing, waitGuard);
                if (cachedPlanValidator.test(sharedPlan)) {
                    return new LookupResult(sharedPlan, true);
                }
                synchronized (entries) {
                    removeEntryIfNotNewer(key, sharedPlan, generation);
                }
                awaitInFlightRemoval(key, existing, waitGuard);
                continue;
            }
            try {
                waitGuard.ensureAllowed();
                FederationSqlPlan plan = compiler.get();
                waitGuard.ensureAllowed();
                long weight = estimateWeight(plan);
                boolean cacheable = weight <= maximumWeightBytes
                    && cachedPlanValidator.test(plan);
                synchronized (entries) {
                    // 缓存写入、成功发布与 close 共享线性化点，避免关闭后的迟到发布。
                    ensureOpen();
                    if (cacheable) {
                        if (putEntry(key, plan, weight, nanoTime.getAsLong(), generation)) {
                            evictToBounds();
                        }
                    }
                    ownFuture.complete(plan);
                }
                return new LookupResult(plan, false);
            } catch (Throwable throwable) {
                ownFuture.completeExceptionally(throwable);
                throw throwable;
            } finally {
                inFlight.remove(key, ownFuture);
                compileSlots.release();
            }
        }
    }

    private void awaitInFlightRemoval(
        PlanCacheKey key,
        CompletableFuture<FederationSqlPlan> completed,
        WaitGuard waitGuard
    ) {
        while (inFlight.get(key) == completed) {
            ensureOpen();
            waitGuard.ensureAllowed();
            LockSupport.parkNanos(Math.min(
                boundedWaitNanos(waitGuard.remainingNanos()),
                TimeUnit.MILLISECONDS.toNanos(1)
            ));
            if (Thread.interrupted()) {
                Thread.currentThread().interrupt();
                throw new FederationSqlException(
                    FederationSqlErrorCode.SQL_COMPILE_FAILED,
                    "waiting for a stale in-flight SQL compilation was interrupted"
                );
            }
        }
    }

    /**
     * 计划缓存查询结果。
     *
     * @param plan 计划
     * @param cacheHit 是否复用缓存或同键 single-flight
     */
    record LookupResult(FederationSqlPlan plan, boolean cacheHit) {
    }

    private static int defaultCompilationConcurrency() {
        return Math.max(1, Math.min(8, Runtime.getRuntime().availableProcessors()));
    }

    private static long defaultMaximumWeight(int maximumEntries) {
        return Math.max(16L * 1024L * 1024L, Math.min(
            256L * 1024L * 1024L,
            maximumEntries * 128L * 1024L
        ));
    }

    private void acquireCompileSlot(WaitGuard waitGuard) {
        while (true) {
            ensureOpen();
            waitGuard.ensureAllowed();
            long waitNanos = boundedWaitNanos(waitGuard.remainingNanos());
            try {
                if (compileSlots.tryAcquire(waitNanos, TimeUnit.NANOSECONDS)) {
                    try {
                        ensureOpen();
                    } catch (RuntimeException exception) {
                        compileSlots.release();
                        throw exception;
                    }
                    return;
                }
            } catch (InterruptedException exception) {
                Thread.currentThread().interrupt();
                throw new FederationSqlException(
                    FederationSqlErrorCode.SQL_COMPILE_FAILED,
                    "waiting for a SQL compile slot was interrupted",
                    exception
                );
            }
        }
    }

    private void ensureOpen() {
        if (closed.get()) {
            throw new FederationSqlException(
                FederationSqlErrorCode.ENGINE_CLOSED,
                "plan cache is closed"
            );
        }
    }

    private static FederationSqlPlan await(
        CompletableFuture<FederationSqlPlan> future,
        WaitGuard waitGuard
    ) {
        while (true) {
            waitGuard.ensureAllowed();
            try {
                return future.get(
                    boundedWaitNanos(waitGuard.remainingNanos()),
                    TimeUnit.NANOSECONDS
                );
            } catch (TimeoutException ignored) {
                // 短轮询使同键等待能及时响应取消和统一截止时间。
            } catch (InterruptedException exception) {
                Thread.currentThread().interrupt();
                throw new FederationSqlException(
                    FederationSqlErrorCode.SQL_COMPILE_FAILED,
                    "waiting for an in-flight SQL compilation was interrupted",
                    exception
                );
            } catch (ExecutionException exception) {
                if (exception.getCause() instanceof RuntimeException runtimeException) {
                    throw runtimeException;
                }
                throw new CompletionException(exception.getCause());
            }
        }
    }

    private static long boundedWaitNanos(long remainingNanos) {
        return Math.max(
            1L,
            Math.min(remainingNanos, TimeUnit.MILLISECONDS.toNanos(50))
        );
    }

    /**
     * 冷编译等待检查器。
     */
    interface WaitGuard {

        /** 检查等待是否仍允许继续。 */
        void ensureAllowed();

        /**
         * 返回剩余等待时间。
         *
         * @return 剩余纳秒
         */
        long remainingNanos();

        /**
         * 返回无限制检查器。
         *
         * @return 无限制检查器
         */
        static WaitGuard unbounded() {
            return UnboundedHolder.INSTANCE;
        }

        /** 无限制实例持有者。 */
        final class UnboundedHolder {
            private static final WaitGuard INSTANCE = new WaitGuard() {
                @Override
                public void ensureAllowed() {
                }

                @Override
                public long remainingNanos() {
                    return Long.MAX_VALUE;
                }
            };

            private UnboundedHolder() {
            }
        }
    }

    /**
     * 返回当前缓存条目数，供测试与指标桥接读取。
     *
     * @return 条目数
     */
    int size() {
        synchronized (entries) {
            removeExpired(nanoTime.getAsLong());
            return entries.size();
        }
    }

    /**
     * 返回当前缓存估算权重，供测试和监控读取。
     *
     * @return 当前估算权重
     */
    long weightBytes() {
        synchronized (entries) {
            removeExpired(nanoTime.getAsLong());
            return currentWeightBytes;
        }
    }

    /**
     * 返回当前可用冷编译槽，供关闭竞态测试确认许可没有泄漏。
     *
     * @return 可用冷编译槽
     */
    int availableCompileSlots() {
        return compileSlots.availablePermits();
    }

    /**
     * 精确移除引用已关闭 Runtime 身份的计划。
     *
     * @param runtime 已关闭 Runtime
     */
    void invalidateRuntime(SourceRuntime runtime) {
        invalidateRuntimeIdentity(
            runtime.definition().sourceId(),
            runtime.definition().revision(),
            runtime.sourceChecksum(),
            runtime.runtimeFingerprint()
        );
    }

    /**
     * 精确移除引用指定 Runtime 身份的计划，供生命周期回调与回归测试复用。
     *
     * @param sourceId 物理源标识
     * @param revision Definition 版本
     * @param checksum Definition 校验和
     * @param runtimeFingerprint 数据库与驱动指纹
     */
    void invalidateRuntimeIdentity(
        com.easyagents.federation.sql.source.SourceId sourceId,
        long revision,
        String checksum,
        String runtimeFingerprint
    ) {
        synchronized (entries) {
            Iterator<Map.Entry<PlanCacheKey, CacheEntry>> iterator = entries.entrySet().iterator();
            while (iterator.hasNext()) {
                CacheEntry entry = iterator.next().getValue();
                boolean matches = entry.plan().sourceRuntimeIdentities().stream().anyMatch(identity ->
                    identity.sourceId().equals(sourceId)
                        && identity.sourceRevision() == revision
                        && identity.sourceChecksum().equals(checksum)
                        && identity.runtimeFingerprint().equals(runtimeFingerprint)
                );
                if (matches) {
                    currentWeightBytes -= entry.weightBytes();
                    iterator.remove();
                }
            }
        }
    }

    /**
     * 清空计划和进行中的编译记录。
     */
    @Override
    public void close() {
        synchronized (entries) {
            if (!closed.compareAndSet(false, true)) {
                return;
            }
            entries.clear();
            currentWeightBytes = 0;
            FederationSqlException failure = new FederationSqlException(
                FederationSqlErrorCode.ENGINE_CLOSED,
                "plan cache was closed during compilation"
            );
            inFlight.values().forEach(future -> future.completeExceptionally(failure));
            inFlight.clear();
        }
    }

    private CacheEntry cachedEntry(PlanCacheKey key, long now) {
        CacheEntry cached = entries.get(key);
        if (cached == null) {
            return null;
        }
        if (isExpired(cached, now)) {
            removeEntryIfSame(key, cached.plan());
            return null;
        }
        return cached;
    }

    private void removeEntryIfSame(PlanCacheKey key, FederationSqlPlan plan) {
        CacheEntry current = entries.get(key);
        if (current != null && current.plan() == plan) {
            entries.remove(key);
            currentWeightBytes -= current.weightBytes();
        }
    }

    private void removeEntryIfNotNewer(
        PlanCacheKey key,
        FederationSqlPlan plan,
        Instant contextGeneration
    ) {
        CacheEntry current = entries.get(key);
        if (current != null
            && current.plan() == plan
            && !current.contextGeneration().isAfter(contextGeneration)) {
            entries.remove(key);
            currentWeightBytes -= current.weightBytes();
        }
    }

    private boolean putEntry(
        PlanCacheKey key,
        FederationSqlPlan plan,
        long weightBytes,
        long now,
        Instant contextGeneration
    ) {
        long statisticsTtlNanos = remainingStatisticsTtlNanos(
            plan.statisticsValidUntil(),
            Instant.now()
        );
        long effectiveTtlNanos = Math.min(ttlNanos, statisticsTtlNanos);
        if (effectiveTtlNanos <= 0L) {
            return false;
        }
        CacheEntry current = entries.get(key);
        if (current != null && current.contextGeneration().isAfter(contextGeneration)) {
            return false;
        }
        CacheEntry previous = entries.put(
            key,
            new CacheEntry(
                plan,
                weightBytes,
                now,
                effectiveTtlNanos,
                contextGeneration
            )
        );
        if (previous != null) {
            currentWeightBytes -= previous.weightBytes();
        }
        currentWeightBytes += weightBytes;
        return true;
    }

    private void evictToBounds() {
        Iterator<Map.Entry<PlanCacheKey, CacheEntry>> iterator = entries.entrySet().iterator();
        while ((entries.size() > maximumEntries || currentWeightBytes > maximumWeightBytes)
            && iterator.hasNext()) {
            CacheEntry eldest = iterator.next().getValue();
            currentWeightBytes -= eldest.weightBytes();
            iterator.remove();
        }
    }

    private void removeExpired(long now) {
        Iterator<Map.Entry<PlanCacheKey, CacheEntry>> iterator = entries.entrySet().iterator();
        while (iterator.hasNext()) {
            CacheEntry entry = iterator.next().getValue();
            if (isExpired(entry, now)) {
                currentWeightBytes -= entry.weightBytes();
                iterator.remove();
            }
        }
    }

    private static long estimateWeight(FederationSqlPlan plan) {
        long[] nodes = {0};
        if (plan.relRoot() != null && plan.relRoot().rel != null) {
            new RelVisitor() {
                @Override
                public void visit(RelNode node, int ordinal, RelNode parent) {
                    nodes[0]++;
                    super.visit(node, ordinal, parent);
                }
            }.go(plan.relRoot().rel);
        }
        long characters = plan.normalizedSql().length() + plan.executableSql().length();
        for (var fragment : plan.fragments()) {
            characters += fragment.executableSql().length();
        }
        return 4_096L
            + nodes[0] * 1_024L
            + characters * 2L
            + plan.columns().size() * 256L
            + plan.fragments().size() * 512L;
    }

    private static long saturatingNanos(Duration duration) {
        try {
            return duration.toNanos();
        } catch (ArithmeticException ignored) {
            return Long.MAX_VALUE;
        }
    }

    private static long remainingStatisticsTtlNanos(Instant validUntil, Instant now) {
        if (validUntil == null || Instant.MAX.equals(validUntil)) {
            return Long.MAX_VALUE;
        }
        if (!validUntil.isAfter(now)) {
            return 0L;
        }
        return saturatingNanos(Duration.between(now, validUntil));
    }

    private static boolean isExpired(CacheEntry entry, long now) {
        return entry.ttlNanos() != Long.MAX_VALUE
            && now - entry.createdAtNanos() >= entry.ttlNanos();
    }

    private record CacheEntry(
        FederationSqlPlan plan,
        long weightBytes,
        long createdAtNanos,
        long ttlNanos,
        Instant contextGeneration
    ) {
    }
}
