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.api.FederationCleanupMetrics;
import com.easyagents.federation.sql.execute.QueryId;
import com.easyagents.federation.sql.execute.StatementLifecycle;
import java.sql.SQLException;
import java.sql.Statement;
import java.time.Duration;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 从 QueryId 建立到 Statement、Cursor 结束的节点本地取消状态。
 */
final class QueryCancellationRegistry implements AutoCloseable {

    private static final System.Logger LOGGER = System.getLogger(
        QueryCancellationRegistry.class.getName()
    );
    private static final int MAX_PENDING_CANCELLATIONS = 4_096;
    private static final int MAX_DEFERRED_CLEANUPS = 4_096;
    private static final long DEFERRED_RETRY_DELAY_MILLIS = 10L;
    private static final long PENDING_CANCELLATION_TTL_NANOS =
        Duration.ofMinutes(10).toNanos();

    private final Map<QueryId, QueryState> queries = new ConcurrentHashMap<>();
    private final Object lifecycleLock = new Object();
    private final LinkedHashMap<QueryId, Long> pendingCancellations =
        new LinkedHashMap<>();
    private final AtomicBoolean closed = new AtomicBoolean();
    private final AtomicBoolean acceptingDeferredCleanups = new AtomicBoolean(true);
    private final AtomicBoolean retryScheduled = new AtomicBoolean();
    private final AtomicLong overflowFallbacks = new AtomicLong();
    private final AtomicLong deferredRetries = new AtomicLong();
    private final AtomicLong unscheduledCleanups = new AtomicLong();
    private final Object deferredLock = new Object();
    private final ArrayBlockingQueue<DeferredCleanup> deferredCleanups =
        new ArrayBlockingQueue<>(MAX_DEFERRED_CLEANUPS);
    private final ThreadPoolExecutor terminationExecutor;
    private final ThreadPoolExecutor cleanupExecutor;
    private final ThreadPoolExecutor overflowExecutor;
    private final ScheduledThreadPoolExecutor retryExecutor;

    /**
     * 创建使用独立 Statement 终止与游标清理队列的取消登记器。
     */
    QueryCancellationRegistry() {
        this(
            executor("easy-agents-federation-jdbc-cancel"),
            executor("easy-agents-federation-cursor-cleanup"),
            executor("easy-agents-federation-resource-overflow")
        );
    }

    /**
     * 创建使用指定执行器的取消登记器，供确定性并发测试使用。
     *
     * @param terminationExecutor Statement 终止执行器
     * @param cleanupExecutor 游标清理执行器
     */
    QueryCancellationRegistry(
        ThreadPoolExecutor terminationExecutor,
        ThreadPoolExecutor cleanupExecutor
    ) {
        this(
            terminationExecutor,
            cleanupExecutor,
            executor("easy-agents-federation-resource-overflow")
        );
    }

    /**
     * 创建使用指定主执行器和过载隔离执行器的取消登记器。
     *
     * @param terminationExecutor Statement 终止执行器
     * @param cleanupExecutor 游标清理执行器
     * @param overflowExecutor 主队列饱和后的独立隔离执行器
     */
    QueryCancellationRegistry(
        ThreadPoolExecutor terminationExecutor,
        ThreadPoolExecutor cleanupExecutor,
        ThreadPoolExecutor overflowExecutor
    ) {
        this.terminationExecutor = Objects.requireNonNull(
            terminationExecutor,
            "terminationExecutor must not be null"
        );
        this.cleanupExecutor = Objects.requireNonNull(
            cleanupExecutor,
            "cleanupExecutor must not be null"
        );
        this.overflowExecutor = Objects.requireNonNull(
            overflowExecutor,
            "overflowExecutor must not be null"
        );
        this.retryExecutor = retryExecutor();
    }

    /**
     * 在查询进入编译、准入或执行前登记取消状态。
     *
     * @param queryId 查询标识
     * @return 查询级登记
     */
    QueryRegistration begin(QueryId queryId) {
        QueryState state;
        synchronized (lifecycleLock) {
            ensureOpen();
            if (consumePendingCancellation(queryId)) {
                throw cancelled(queryId, null);
            }
            state = new QueryState(queryId);
            QueryState previous = queries.putIfAbsent(queryId, state);
            if (previous != null) {
                throw new FederationSqlException(
                    FederationSqlErrorCode.EXECUTION_FAILED,
                    "query id is already active: " + queryId.value()
                );
            }
        }
        return new QueryRegistration(queryId, state);
    }

    /**
     * 取消指定查询。
     *
     * @param queryId 查询标识
     * @return 是否找到活动查询；未找到时仍保留有界的短期预取消状态
     */
    boolean cancel(QueryId queryId) {
        return cancelOutcome(queryId).found();
    }

    /**
     * 取消指定查询并返回原子确定的终止原因。
     *
     * @param queryId 查询标识
     * @return 取消观察结果
     */
    CancellationOutcome cancelOutcome(QueryId queryId) {
        QueryState state;
        synchronized (lifecycleLock) {
            state = queries.get(queryId);
            if (state == null) {
                rememberPendingCancellation(queryId);
                return new CancellationOutcome(false, TerminationReason.NONE);
            }
        }
        return new CancellationOutcome(true, state.requestCancellation());
    }

    /**
     * 在隔离线程中执行可能被 JDBC 驱动阻塞的游标清理。
     *
     * @param cleanup 清理动作
     * @return 是否成功提交
     */
    boolean submitCleanup(Runnable cleanup) {
        if (cleanup == null) {
            return false;
        }
        return submit(
            cleanupExecutor,
            cleanup,
            "query cursor cleanup failed",
            "cursor cleanup"
        );
    }

    /**
     * 返回主队列饱和后转交隔离执行器的保护性降级次数。
     *
     * @return 隔离降级次数
     */
    long overflowFallbackCount() {
        return overflowFallbacks.get();
    }

    /**
     * 返回所有隔离队列均饱和后未能调度的清理次数。
     *
     * @return 未调度清理次数
     */
    long unscheduledCleanupCount() {
        return unscheduledCleanups.get();
    }

    /**
     * 返回清理通道的不可变观测快照。
     *
     * @return 清理通道指标
     */
    FederationCleanupMetrics cleanupMetrics() {
        return new FederationCleanupMetrics(
            overflowFallbacks.get(),
            deferredRetries.get(),
            unscheduledCleanups.get(),
            deferredCleanups.size()
        );
    }

    /**
     * 尝试取消当前节点所有查询，并保留第一个错误。
     */
    @Override
    public void close() {
        synchronized (lifecycleLock) {
            if (!closed.compareAndSet(false, true)) {
                return;
            }
            pendingCancellations.clear();
        }
        for (QueryState state : queries.values()) {
            state.requestCancellation();
        }
        transferQueuedOnShutdown(terminationExecutor, "JDBC statement termination");
        transferQueuedOnShutdown(cleanupExecutor, "cursor cleanup");
        drainDeferredCleanups();
        acceptingDeferredCleanups.set(false);
        retryExecutor.shutdownNow();
        shutdownOverflow();
    }

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

    private boolean consumePendingCancellation(QueryId queryId) {
        cleanupPendingCancellations(System.nanoTime());
        return pendingCancellations.remove(queryId) != null;
    }

    private void rememberPendingCancellation(QueryId queryId) {
        long now = System.nanoTime();
        cleanupPendingCancellations(now);
        pendingCancellations.put(
            queryId,
            now + PENDING_CANCELLATION_TTL_NANOS
        );
        while (pendingCancellations.size() > MAX_PENDING_CANCELLATIONS) {
            Iterator<QueryId> iterator = pendingCancellations.keySet().iterator();
            iterator.next();
            iterator.remove();
        }
    }

    private void cleanupPendingCancellations(long now) {
        pendingCancellations.entrySet().removeIf(
            entry -> entry.getValue() - now <= 0L
        );
    }

    /**
     * 单次查询的取消登记，直到 Cursor 和 Core 资源完成关闭后才清理终态。
     */
    final class QueryRegistration implements AutoCloseable {

        private final QueryId queryId;
        private final QueryState state;
        private final AtomicBoolean released = new AtomicBoolean();

        private QueryRegistration(QueryId queryId, QueryState state) {
            this.queryId = queryId;
            this.state = state;
        }

        /**
         * 返回 Adapter 用于登记 Statement 的生命周期回调。
         *
         * @return Statement 生命周期
         */
        StatementLifecycle statementLifecycle() {
            return state.statementLifecycle();
        }

        /**
         * 返回查询是否已经收到取消请求。
         *
         * @return 是否已取消
         */
        boolean cancellationRequested() {
            return state.cancellationRequested();
        }

        /**
         * 已取消时立即抛出稳定错误。
         */
        void ensureNotCancelled() {
            state.ensureActive();
        }

        /**
         * 将查询推进为超时终态，并关闭当前已登记的全部 Statement。
         */
        void requestTimeout() {
            state.requestTimeout();
        }

        /**
         * 将查询推进为超时并返回原子确定的终止原因。
         *
         * @return 已确定的终止原因
         */
        TerminationReason requestTimeoutOutcome() {
            return state.requestTimeout();
        }

        /**
         * 返回查询是否已经超过统一执行时限。
         *
         * @return 是否已超时
         */
        boolean timeoutRequested() {
            return state.timeoutRequested();
        }

        /**
         * 清理活动 QueryId；状态对象仍供并发 JDBC 消费线程判断取消原因。
         */
        @Override
        public void close() {
            if (released.compareAndSet(false, true)) {
                queries.remove(queryId, state);
                state.releaseStatementReference();
            }
        }
    }

    private final class QueryState {

        private final QueryId queryId;
        private final Map<Statement, StatementTermination> statements = new ConcurrentHashMap<>();
        private final AtomicReference<TerminationReason> terminationReason =
            new AtomicReference<>(TerminationReason.NONE);
        private final StatementLifecycle statementLifecycle = new StatementLifecycle() {
            @Override
            public void register(Statement candidate) {
                if (candidate == null) {
                    throw new IllegalArgumentException("statement must not be null");
                }
                ensureActive();
                StatementTermination termination = new StatementTermination(candidate);
                StatementTermination previous = statements.putIfAbsent(candidate, termination);
                if (previous != null && previous.statement != candidate) {
                    throw new FederationSqlException(
                        FederationSqlErrorCode.EXECUTION_FAILED,
                        "query id already registered a different statement instance"
                    );
                }
                if (terminationRequested()) {
                    StatementTermination active = previous == null ? termination : previous;
                    submitTermination(
                        () -> active.terminate(queryId),
                        "late JDBC statement termination failed"
                    );
                    ensureActive();
                }
            }

            @Override
            public void unregister(Statement candidate) {
                statements.remove(candidate);
            }

            @Override
            public boolean cancellationRequested() {
                return QueryState.this.terminationRequested();
            }

            @Override
            public boolean timeoutRequested() {
                return QueryState.this.timeoutRequested();
            }
        };

        private QueryState(QueryId queryId) {
            this.queryId = queryId;
        }

        private StatementLifecycle statementLifecycle() {
            return statementLifecycle;
        }

        private boolean cancellationRequested() {
            return terminationRequested();
        }

        private boolean terminationRequested() {
            return terminationReason.get() != TerminationReason.NONE;
        }

        private boolean timeoutRequested() {
            return terminationReason.get() == TerminationReason.TIMED_OUT;
        }

        private void ensureActive() {
            TerminationReason reason = terminationReason.get();
            if (reason == TerminationReason.TIMED_OUT) {
                throw timedOut(queryId, null);
            }
            if (reason == TerminationReason.CANCELLED) {
                throw cancelled(queryId, null);
            }
        }

        private TerminationReason requestCancellation() {
            if (terminationReason.compareAndSet(
                TerminationReason.NONE,
                TerminationReason.CANCELLED
            )) {
                terminateStatements();
            }
            return terminationReason.get();
        }

        private TerminationReason requestTimeout() {
            if (terminationReason.compareAndSet(
                TerminationReason.NONE,
                TerminationReason.TIMED_OUT
            )) {
                terminateStatements();
            }
            return terminationReason.get();
        }

        private void terminateStatements() {
            for (StatementTermination termination : List.copyOf(statements.values())) {
                submitTermination(
                    () -> termination.terminate(queryId),
                    "JDBC statement termination failed"
                );
            }
        }

        private void releaseStatementReference() {
            statements.clear();
        }
    }

    enum TerminationReason {
        NONE,
        CANCELLED,
        TIMED_OUT
    }

    /**
     * 一次取消请求观察到的活动状态和原子终态。
     *
     * @param found 是否找到活动查询
     * @param reason 已确定的终止原因
     */
    record CancellationOutcome(boolean found, TerminationReason reason) {
    }

    /**
     * 每个 Statement 独立维护取消和关闭幂等状态。
     */
    private static final class StatementTermination {

        private final Statement statement;
        private final AtomicBoolean cancelIssued = new AtomicBoolean();
        private final AtomicBoolean closeIssued = new AtomicBoolean();

        private StatementTermination(Statement statement) {
            this.statement = statement;
        }

        private void terminate(QueryId queryId) {
            SQLException failure = null;
            if (cancelIssued.compareAndSet(false, true)) {
                try {
                    statement.cancel();
                } catch (SQLException exception) {
                    cancelIssued.set(false);
                    failure = exception;
                }
            }
            // close() 同时覆盖 register 与 executeQuery 之间的 JDBC 取消空窗。
            if (closeIssued.compareAndSet(false, true)) {
                try {
                    statement.close();
                } catch (SQLException exception) {
                    closeIssued.set(false);
                    if (failure == null) {
                        failure = exception;
                    } else {
                        failure.addSuppressed(exception);
                    }
                }
            }
            if (failure != null) {
                throw new FederationSqlException(
                    FederationSqlErrorCode.EXECUTION_FAILED,
                    "failed to cancel query " + queryId.value(),
                    failure
                );
            }
        }
    }

    private static FederationSqlException cancelled(QueryId queryId, Throwable cause) {
        String message = "query was cancelled: " + queryId.value();
        return cause == null
            ? new FederationSqlException(FederationSqlErrorCode.QUERY_CANCELLED, message)
            : new FederationSqlException(FederationSqlErrorCode.QUERY_CANCELLED, message, cause);
    }

    private static FederationSqlException timedOut(QueryId queryId, Throwable cause) {
        String message = "query execution timed out: " + queryId.value();
        return cause == null
            ? new FederationSqlException(FederationSqlErrorCode.QUERY_TIMEOUT, message)
            : new FederationSqlException(FederationSqlErrorCode.QUERY_TIMEOUT, message, cause);
    }

    private boolean submitTermination(Runnable action, String failureMessage) {
        return submit(
            terminationExecutor,
            action,
            failureMessage,
            "JDBC statement termination"
        );
    }

    private boolean submit(
        ThreadPoolExecutor executor,
        Runnable action,
        String failureMessage,
        String actionLabel
    ) {
        Runnable guarded = () -> {
            try {
                runGuarded(action, failureMessage);
            } finally {
                drainDeferredCleanups();
            }
        };
        try {
            executor.execute(guarded);
            return true;
        } catch (RejectedExecutionException exception) {
            overflowFallbacks.incrementAndGet();
            LOGGER.log(
                System.Logger.Level.WARNING,
                actionLabel
                    + " queue is unavailable; handing cleanup to the overflow isolator: "
                    + exception.getMessage()
            );
            submitOverflow(guarded, actionLabel);
            return false;
        }
    }

    private void submitOverflow(Runnable guarded, String actionLabel) {
        try {
            overflowExecutor.execute(guarded);
        } catch (RejectedExecutionException overflow) {
            deferCleanup(new DeferredCleanup(guarded, actionLabel), overflow);
        }
    }

    private void deferCleanup(DeferredCleanup cleanup, RejectedExecutionException cause) {
        boolean accepted;
        synchronized (deferredLock) {
            accepted = acceptingDeferredCleanups.get() && deferredCleanups.offer(cleanup);
        }
        if (!accepted) {
            recordUnresolved(
                1L,
                cleanup.actionLabel()
                    + " could not be retained because all bounded cleanup channels are saturated",
                cause
            );
            return;
        }
        deferredRetries.incrementAndGet();
        LOGGER.log(
            System.Logger.Level.WARNING,
            cleanup.actionLabel()
                + " entered the bounded deferred cleanup queue after both executors rejected it"
        );
        scheduleDeferredRetry();
    }

    private void scheduleDeferredRetry() {
        if (!acceptingDeferredCleanups.get()
            || !retryScheduled.compareAndSet(false, true)) {
            return;
        }
        try {
            retryExecutor.schedule(
                this::retryDeferredCleanups,
                DEFERRED_RETRY_DELAY_MILLIS,
                TimeUnit.MILLISECONDS
            );
        } catch (RejectedExecutionException exception) {
            retryScheduled.set(false);
            if (acceptingDeferredCleanups.get()) {
                LOGGER.log(
                    System.Logger.Level.ERROR,
                    "deferred JDBC cleanup retry scheduler is unavailable",
                    exception
                );
            }
        }
    }

    private void retryDeferredCleanups() {
        retryScheduled.set(false);
        drainDeferredCleanups();
        if (!deferredCleanups.isEmpty()) {
            scheduleDeferredRetry();
        }
    }

    private void drainDeferredCleanups() {
        while (acceptingDeferredCleanups.get()) {
            DeferredCleanup cleanup;
            synchronized (deferredLock) {
                cleanup = deferredCleanups.poll();
            }
            if (cleanup == null) {
                return;
            }
            try {
                overflowExecutor.execute(cleanup.action());
            } catch (RejectedExecutionException exception) {
                boolean restored;
                synchronized (deferredLock) {
                    restored = acceptingDeferredCleanups.get()
                        && deferredCleanups.offer(cleanup);
                }
                if (!restored) {
                    recordUnresolved(
                        1L,
                        cleanup.actionLabel()
                            + " could not be retained while retrying deferred cleanup",
                        exception
                    );
                } else {
                    scheduleDeferredRetry();
                }
                return;
            }
        }
    }

    private static void runGuarded(Runnable action, String failureMessage) {
        try {
            action.run();
        } catch (RuntimeException exception) {
            LOGGER.log(System.Logger.Level.ERROR, failureMessage, exception);
        }
    }

    private void transferQueuedOnShutdown(
        ThreadPoolExecutor executor,
        String actionLabel
    ) {
        executor.shutdown();
        boolean interrupted = false;
        List<Runnable> queued = List.of();
        boolean terminated = false;
        try {
            if (!executor.awaitTermination(200, TimeUnit.MILLISECONDS)) {
                queued = executor.shutdownNow();
            } else {
                terminated = true;
            }
        } catch (InterruptedException exception) {
            interrupted = true;
            queued = executor.shutdownNow();
        }
        // 关闭线程只做有界等待；尚未开始的阻塞清理转交隔离执行器。
        queued.forEach(task -> submitOverflow(task, actionLabel));
        if (!terminated && executor.getActiveCount() > 0) {
            recordUnresolved(
                executor.getActiveCount(),
                actionLabel + " remained active after the bounded engine shutdown window",
                null
            );
        }
        if (interrupted) {
            Thread.currentThread().interrupt();
        }
    }

    private void shutdownOverflow() {
        overflowExecutor.shutdown();
        boolean interrupted = false;
        List<Runnable> discarded = List.of();
        try {
            if (!overflowExecutor.awaitTermination(200, TimeUnit.MILLISECONDS)) {
                discarded = overflowExecutor.shutdownNow();
            }
        } catch (InterruptedException exception) {
            interrupted = true;
            discarded = overflowExecutor.shutdownNow();
        }
        int deferred;
        synchronized (deferredLock) {
            deferred = deferredCleanups.size();
            deferredCleanups.clear();
        }
        long unresolved = (long) discarded.size()
            + overflowExecutor.getActiveCount()
            + deferred;
        if (unresolved > 0L) {
            recordUnresolved(
                unresolved,
                "JDBC cleanup tasks remained after the bounded engine shutdown window",
                null
            );
        }
        if (interrupted) {
            Thread.currentThread().interrupt();
        }
    }

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

    private static ScheduledThreadPoolExecutor retryExecutor() {
        ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(
            1,
            runnable -> {
                Thread thread = new Thread(
                    runnable,
                    "easy-agents-federation-resource-retry"
                );
                thread.setDaemon(true);
                return thread;
            },
            new ThreadPoolExecutor.AbortPolicy()
        );
        executor.setExecuteExistingDelayedTasksAfterShutdownPolicy(false);
        executor.setRemoveOnCancelPolicy(true);
        return executor;
    }

    private void recordUnresolved(long count, String message, Throwable cause) {
        unscheduledCleanups.addAndGet(count);
        if (cause == null) {
            LOGGER.log(System.Logger.Level.ERROR, count + " " + message);
        } else {
            LOGGER.log(System.Logger.Level.ERROR, count + " " + message, cause);
        }
    }

    private record DeferredCleanup(Runnable action, String actionLabel) {
    }
}
