package com.easyagents.federation.sql.runtime;

import com.easyagents.federation.sql.adapter.FederationStatisticsCollectionContext;
import com.easyagents.federation.sql.adapter.FederationStatisticsCollector;
import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import com.easyagents.federation.sql.federation.FederationStatisticsSnapshot;
import com.easyagents.federation.sql.federation.FederationTableStatistics;
import com.easyagents.federation.sql.federation.FederationTableStatisticsProvider;
import com.easyagents.federation.sql.source.SourceId;
import java.sql.Connection;
import java.sql.SQLException;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 使用数据库 Adapter 自动采集、缓存并降级联邦查询统计。
 */
final class AdapterFederationStatisticsProvider
    implements FederationTableStatisticsProvider, AutoCloseable {

    private static final Logger LOG = LoggerFactory.getLogger(
        AdapterFederationStatisticsProvider.class
    );
    private static final Duration DEFAULT_TTL = Duration.ofMinutes(30);
    private static final Duration DEFAULT_FAILURE_RETRY_DELAY = Duration.ofMinutes(1);
    private static final int DEFAULT_QUERY_TIMEOUT_SECONDS = 5;
    private static final int MAXIMUM_PARALLEL_REFRESHES = 4;
    private static final int MAXIMUM_PENDING_REFRESHES = 64;

    private final Duration ttl;
    private final Duration failureRetryDelay;
    private final int queryTimeoutSeconds;
    private final ConcurrentMap<SourceId, Object> sourceLocks = new ConcurrentHashMap<>();
    private final ConcurrentMap<SourceRevision, CompletableFuture<Void>> inFlight =
        new ConcurrentHashMap<>();
    private final ExecutorService refreshExecutor;
    private final AtomicBoolean closed = new AtomicBoolean();
    private volatile RegistryState state = RegistryState.empty();
    private long versionSequence;

    /**
     * 使用默认有效期、失败限频和有界并行度创建自动统计 Provider。
     */
    AdapterFederationStatisticsProvider() {
        this(
            DEFAULT_TTL,
            DEFAULT_FAILURE_RETRY_DELAY,
            DEFAULT_QUERY_TIMEOUT_SECONDS,
            Math.min(
                MAXIMUM_PARALLEL_REFRESHES,
                Math.max(1, Runtime.getRuntime().availableProcessors())
            )
        );
    }

    /**
     * 创建可测试的自动统计 Provider。
     *
     * @param ttl 统计有效期
     * @param failureRetryDelay 失败后的最短重试间隔
     * @param queryTimeoutSeconds 单条目录查询超时
     * @param parallelism 不同物理源的最大并行采集数
     */
    AdapterFederationStatisticsProvider(
        Duration ttl,
        Duration failureRetryDelay,
        int queryTimeoutSeconds,
        int parallelism
    ) {
        this.ttl = positive(ttl, "ttl");
        this.failureRetryDelay = positive(failureRetryDelay, "failureRetryDelay");
        if (queryTimeoutSeconds <= 0) {
            throw new IllegalArgumentException("queryTimeoutSeconds must be positive");
        }
        if (parallelism <= 0) {
            throw new IllegalArgumentException("parallelism must be positive");
        }
        this.queryTimeoutSeconds = queryTimeoutSeconds;
        this.refreshExecutor = new ThreadPoolExecutor(
            parallelism,
            parallelism,
            0L,
            TimeUnit.MILLISECONDS,
            new ArrayBlockingQueue<>(MAXIMUM_PENDING_REFRESHES),
            daemonThreadFactory()
        );
    }

    /**
     * 返回当前冻结统计快照。
     *
     * @return 无数据库访问副作用的不可变快照
     */
    @Override
    public FederationStatisticsSnapshot snapshot() {
        return state.snapshot();
    }

    /**
     * 异步刷新查询实际涉及的物理源统计。
     *
     * <p>同一 source revision 的并发刷新合并为一个任务，不同物理源最多按固定
     * 并行度同时读取目录。该方法不阻塞调用线程；需要等待统计的 Explain 调用可
     * 在自身截止时间内等待返回的 Future。失败时保留旧统计并限频。</p>
     *
     * @param querySnapshot 已持有 Runtime lease 的查询范围快照
     * @return 所有已安排刷新完成时结束的 Future
     */
    CompletableFuture<Void> refreshIfNeeded(FederationQueryScopeSnapshot querySnapshot) {
        if (closed.get()) {
            return CompletableFuture.completedFuture(null);
        }
        Map<SourceId, SourceRuntime> runtimes = new LinkedHashMap<>();
        querySnapshot.runtimesByBinding().values().forEach(runtime ->
            runtimes.putIfAbsent(runtime.definition().sourceId(), runtime)
        );
        Instant now = Instant.now();
        List<CompletableFuture<Void>> refreshes = new ArrayList<>();
        for (SourceRuntime runtime : runtimes.values()) {
            if (!fresh(state.sources().get(runtime.definition().sourceId()), runtime, now)) {
                refreshes.add(refreshAsync(runtime));
            }
        }
        return refreshes.isEmpty()
            ? CompletableFuture.completedFuture(null)
            : CompletableFuture.allOf(refreshes.toArray(CompletableFuture[]::new));
    }

    /**
     * 在 Runtime 最终关闭后清除同 revision 的节点本地统计。
     *
     * @param runtime 已关闭 Runtime
     */
    void invalidate(SourceRuntime runtime) {
        if (closed.get()) {
            return;
        }
        SourceId sourceId = runtime.definition().sourceId();
        Object lock = sourceLocks.computeIfAbsent(sourceId, ignored -> new Object());
        synchronized (lock) {
            SourceStatistics current = state.sources().get(sourceId);
            if (current != null && current.revision() == runtime.definition().revision()) {
                replace(sourceId, null);
            }
        }
        // 锁对象必须覆盖数据源的完整生命周期，避免旧 revision 关闭与新 revision
        // 刷新交错时为同一数据源创建两个临界区。锁会在 Provider 关闭时统一释放。
    }

    /**
     * 关闭后台采集器并释放节点本地统计。
     */
    @Override
    public void close() {
        if (!closed.compareAndSet(false, true)) {
            return;
        }
        refreshExecutor.shutdownNow();
        try {
            if (!refreshExecutor.awaitTermination(
                queryTimeoutSeconds + 1L,
                TimeUnit.SECONDS
            )) {
                LOG.warn(
                    "federation statistics refresh did not stop within the bounded shutdown window"
                );
            }
        } catch (InterruptedException exception) {
            Thread.currentThread().interrupt();
            LOG.warn("interrupted while closing federation statistics refresh", exception);
        }
        inFlight.clear();
        sourceLocks.clear();
        synchronized (this) {
            state = RegistryState.empty();
        }
    }

    /**
     * 合并同一 source revision 的并发刷新。
     *
     * @param runtime 待刷新 Runtime
     * @return 可等待的共享刷新任务
     */
    private CompletableFuture<Void> refreshAsync(SourceRuntime runtime) {
        SourceRevision key = new SourceRevision(
            runtime.definition().sourceId(),
            runtime.definition().revision()
        );
        CompletableFuture<Void> created = new CompletableFuture<>();
        CompletableFuture<Void> existing = inFlight.putIfAbsent(key, created);
        if (existing != null) {
            return existing;
        }
        Runnable task = () -> {
            try {
                try (SourceRuntime.RuntimeLease lease = runtime.acquire()) {
                    // 后台采集必须持有独立租约，避免查询快照释放后旧 Runtime
                    // 在目录连接仍被使用时关闭连接池。
                    refresh(lease.runtime());
                }
                // 租约释放可能触发旧 Runtime 关闭，完成信号必须晚于资源收口。
                created.complete(null);
            } catch (FederationSqlException exception) {
                if (exception.errorCode() == FederationSqlErrorCode.SOURCE_REVISION_NOT_READY) {
                    // 已退休 revision 的统计没有继续采集的价值，保留现有估算即可。
                    created.complete(null);
                } else {
                    created.completeExceptionally(exception);
                }
            } catch (RuntimeException exception) {
                created.completeExceptionally(exception);
            } finally {
                inFlight.remove(key, created);
            }
        };
        try {
            refreshExecutor.execute(task);
        } catch (RejectedExecutionException exception) {
            // 查询线程不能在统计线程池饱和时退化为同步目录查询。
            created.complete(null);
            inFlight.remove(key, created);
            if (!closed.get()) {
                LOG.warn(
                    "federation statistics refresh rejected; using existing or default estimates, sourceId={}, revision={}",
                    runtime.definition().sourceId(),
                    runtime.definition().revision()
                );
            }
        }
        return created;
    }

    /**
     * 在源级临界区读取数据库目录并发布统计。
     *
     * @param runtime 数据源 Runtime
     */
    private void refresh(SourceRuntime runtime) {
        SourceId sourceId = runtime.definition().sourceId();
        long revision = runtime.definition().revision();
        Object lock = sourceLocks.computeIfAbsent(sourceId, ignored -> new Object());
        synchronized (lock) {
            Instant now = Instant.now();
            SourceStatistics current = state.sources().get(sourceId);
            if (current != null && current.revision() > revision) {
                return;
            }
            if (fresh(current, runtime, now)) {
                return;
            }
            try {
                Map<FederationStatisticsSnapshot.TableKey, FederationTableStatistics> tables =
                    collect(runtime, now);
                if (closed.get()) {
                    return;
                }
                replace(sourceId, new SourceStatistics(
                    revision,
                    validateSource(sourceId, tables),
                    now.plus(ttl)
                ));
            } catch (SQLException | RuntimeException exception) {
                if (!closed.get()) {
                    throttleAfterFailure(sourceId, revision, now);
                }
                LOG.warn(
                    "federation statistics collection failed; using existing or default estimates, sourceId={}, revision={}",
                    sourceId,
                    revision,
                    exception
                );
            }
        }
    }

    /**
     * 借用运行时连接并调用 Adapter 统计采集器。
     *
     * @param runtime 数据源 Runtime
     * @param collectedAt 采集时间
     * @return 表统计映射
     * @throws SQLException 获取连接或读取目录失败
     */
    private Map<FederationStatisticsSnapshot.TableKey, FederationTableStatistics> collect(
        SourceRuntime runtime,
        Instant collectedAt
    ) throws SQLException {
        FederationStatisticsCollector collector = runtime.adapter().statisticsCollector()
            .orElse(null);
        if (collector == null) {
            return Map.of();
        }
        try (Connection connection = runtime.handle().dataSource().getConnection()) {
            markReadOnly(connection);
            Map<FederationStatisticsSnapshot.TableKey, FederationTableStatistics> collected =
                collector.collect(new FederationStatisticsCollectionContext(
                    runtime.definition(),
                    connection,
                    collectedAt,
                    collectedAt.plus(ttl),
                    queryTimeoutSeconds
                ));
            return Map.copyOf(collected == null ? Map.of() : collected);
        }
    }

    /**
     * 尽力将目录连接标记为只读；驱动不支持时仍由只读 SQL 保证安全。
     *
     * @param connection JDBC 连接
     */
    private void markReadOnly(Connection connection) {
        try {
            if (!connection.isReadOnly()) {
                connection.setReadOnly(true);
            }
        } catch (SQLException exception) {
            LOG.debug("JDBC driver does not support setting statistics connection read-only", exception);
        }
    }

    /**
     * 校验 Adapter 只能发布当前物理源的统计。
     *
     * @param sourceId 当前物理源
     * @param tables Adapter 统计
     * @return 冻结统计映射
     */
    private Map<FederationStatisticsSnapshot.TableKey, FederationTableStatistics> validateSource(
        SourceId sourceId,
        Map<FederationStatisticsSnapshot.TableKey, FederationTableStatistics> tables
    ) {
        for (FederationStatisticsSnapshot.TableKey key : tables.keySet()) {
            if (!sourceId.equals(key.sourceId())) {
                throw new IllegalArgumentException(
                    "statistics collector returned a table for another source: " + key.sourceId()
                );
            }
        }
        return Map.copyOf(tables);
    }

    /**
     * 判断当前 revision 的统计是否仍在刷新窗口内。
     *
     * @param statistics 已缓存源统计
     * @param runtime 当前 Runtime
     * @param now 当前时间
     * @return 可直接使用时为 true
     */
    private boolean fresh(SourceStatistics statistics, SourceRuntime runtime, Instant now) {
        return statistics != null
            && statistics.revision() == runtime.definition().revision()
            && statistics.refreshAfter().isAfter(now);
    }

    /**
     * 采集失败时保留同 revision 旧值并限制后续重试频率。
     *
     * @param sourceId 物理源
     * @param revision 当前 revision
     * @param failedAt 失败时间
     */
    private void throttleAfterFailure(SourceId sourceId, long revision, Instant failedAt) {
        SourceStatistics current = state.sources().get(sourceId);
        Map<FederationStatisticsSnapshot.TableKey, FederationTableStatistics> tables =
            current != null && current.revision() == revision ? current.tables() : Map.of();
        replace(sourceId, new SourceStatistics(
            revision,
            tables,
            failedAt.plus(failureRetryDelay)
        ));
    }

    /**
     * 原子替换单源统计并重建聚合快照。
     *
     * @param sourceId 物理源
     * @param statistics 新统计；null 表示移除
     */
    private synchronized void replace(SourceId sourceId, SourceStatistics statistics) {
        Map<SourceId, SourceStatistics> sources = new LinkedHashMap<>(state.sources());
        if (statistics == null) {
            sources.remove(sourceId);
        } else {
            sources.put(sourceId, statistics);
        }
        Map<FederationStatisticsSnapshot.TableKey, FederationTableStatistics> tables =
            new LinkedHashMap<>();
        sources.values().forEach(value -> tables.putAll(value.tables()));
        versionSequence++;
        state = new RegistryState(
            Map.copyOf(sources),
            new FederationStatisticsSnapshot(
                "adapter-" + versionSequence,
                Instant.now(),
                tables
            )
        );
    }

    /**
     * 校验正 Duration。
     *
     * @param value Duration
     * @param name 参数名
     * @return 原值
     */
    private static Duration positive(Duration value, String name) {
        if (value == null || value.isZero() || value.isNegative()) {
            throw new IllegalArgumentException(name + " must be positive");
        }
        return value;
    }

    /**
     * 创建统计采集守护线程工厂。
     *
     * @return 守护线程工厂
     */
    private static ThreadFactory daemonThreadFactory() {
        AtomicInteger sequence = new AtomicInteger();
        return task -> {
            Thread thread = new Thread(
                task,
                "federation-statistics-" + sequence.incrementAndGet()
            );
            thread.setDaemon(true);
            return thread;
        };
    }

    /**
     * 单物理源统计及刷新边界。
     *
     * @param revision Definition revision
     * @param tables 表统计
     * @param refreshAfter 下次允许刷新时间
     */
    private record SourceStatistics(
        long revision,
        Map<FederationStatisticsSnapshot.TableKey, FederationTableStatistics> tables,
        Instant refreshAfter
    ) {

        /**
         * 冻结源统计。
         */
        private SourceStatistics {
            tables = Map.copyOf(tables == null ? Map.of() : tables);
            refreshAfter = refreshAfter == null ? Instant.EPOCH : refreshAfter;
        }
    }

    /**
     * 正在采集的物理源 revision。
     *
     * @param sourceId 物理源
     * @param revision Definition revision
     */
    private record SourceRevision(SourceId sourceId, long revision) {
    }

    /**
     * 节点当前聚合统计状态。
     *
     * @param sources 按物理源索引的统计
     * @param snapshot 冻结聚合快照
     */
    private record RegistryState(
        Map<SourceId, SourceStatistics> sources,
        FederationStatisticsSnapshot snapshot
    ) {

        /**
         * 创建空统计状态。
         *
         * @return 空状态
         */
        private static RegistryState empty() {
            return new RegistryState(Map.of(), FederationStatisticsSnapshot.empty());
        }

        /**
         * 冻结聚合状态。
         */
        private RegistryState {
            sources = Map.copyOf(sources == null ? Map.of() : sources);
        }
    }
}
