package com.easyagents.federation.sql.adapter.jdbc;

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.execute.FederationExecutionGuard;
import com.easyagents.federation.sql.execute.FederationExecutionObserver;
import com.easyagents.federation.sql.execute.FederationFragmentExecutionContext;
import com.easyagents.federation.sql.execute.QueryId;
import com.easyagents.federation.sql.execute.SqlExecutionOptions;
import com.easyagents.federation.sql.execute.StatementLifecycle;
import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.io.PrintWriter;
import java.io.Reader;
import java.io.StringReader;
import java.lang.reflect.Proxy;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.SQLTimeoutException;
import java.sql.SQLTransientConnectionException;
import java.sql.Statement;
import java.util.List;
import java.util.Map;
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.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.logging.Logger;
import javax.sql.DataSource;
import org.junit.Assert;
import org.junit.Test;

/**
 * JDBC 分片执行器的连接获取错误边界测试。
 */
public class JdbcFederationFragmentExecutorTest {

    /**
     * 验证连接池等待跨过查询截止时间时保留统一查询超时错误。
     */
    @Test
    public void shouldPreferQueryDeadlineOverConnectionPoolTimeout() {
        AtomicInteger checks = new AtomicInteger();
        FederationExecutionGuard guard = new FederationExecutionGuard() {
            @Override
            public void ensureAllowed() {
                if (checks.incrementAndGet() > 1) {
                    throw new FederationSqlException(
                        FederationSqlErrorCode.QUERY_TIMEOUT,
                        "query deadline reached"
                    );
                }
            }

            @Override
            public long remainingNanos() {
                return 1L;
            }
        };

        FederationSqlException failure = expectFailure(context(guard));

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

    /**
     * 验证截止时间仍有效时保留连接获取超时分类。
     */
    @Test
    public void shouldReportConnectionAcquisitionTimeoutBeforeQueryDeadline() {
        FederationSqlException failure = expectFailure(
            context(FederationExecutionGuard.none())
        );

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

    /**
     * 验证显式取消先到达时，驱动的 SQLTimeoutException 不会覆盖取消终态。
     */
    @Test
    public void shouldPreserveCancellationWhenDriverReportsExecutionTimeout() {
        TerminalLifecycle lifecycle = new TerminalLifecycle(false, true);
        AtomicBoolean statementClosed = new AtomicBoolean();
        AtomicBoolean connectionClosed = new AtomicBoolean();

        FederationSqlException failure = expectFailure(context(
            FederationExecutionGuard.none(),
            executionTimeoutDataSource(statementClosed, connectionClosed),
            lifecycle
        ));

        Assert.assertEquals(FederationSqlErrorCode.QUERY_CANCELLED, failure.errorCode());
        Assert.assertTrue(lifecycle.unregistered.get());
        Assert.assertTrue(statementClosed.get());
        Assert.assertTrue(connectionClosed.get());
    }

    /**
     * 验证结果读取阶段同样保留已经先到达的显式取消终态。
     */
    @Test
    public void shouldPreserveCancellationWhenDriverReportsResultTimeout() {
        TerminalLifecycle lifecycle = new TerminalLifecycle(false, true);
        JdbcFederationResultCursor cursor = failingCursor(
            lifecycle,
            new SQLTimeoutException("driver reported timeout after cancel")
        );

        FederationSqlException failure = expectCursorFailure(cursor);

        Assert.assertEquals(FederationSqlErrorCode.QUERY_CANCELLED, failure.errorCode());
        Assert.assertTrue(lifecycle.unregistered.get());
    }

    /**
     * 验证统一超时先到达时，驱动普通异常仍保持超时终态。
     */
    @Test
    public void shouldPreserveTimeoutWhenDriverReportsGenericReadFailure() {
        TerminalLifecycle lifecycle = new TerminalLifecycle(true, true);
        JdbcFederationResultCursor cursor = failingCursor(
            lifecycle,
            new SQLException("statement was closed by timeout task")
        );

        FederationSqlException failure = expectCursorFailure(cursor);

        Assert.assertEquals(FederationSqlErrorCode.QUERY_TIMEOUT, failure.errorCode());
        Assert.assertTrue(lifecycle.unregistered.get());
    }

    /**
     * 验证二进制流取得后发生取消时，后续流读取立即终止并释放 JDBC 资源。
     *
     * @throws Exception 流读取失败
     */
    @Test
    public void shouldStopBinaryStreamReadAfterCancellation() throws Exception {
        AtomicBoolean cancelled = new AtomicBoolean();
        TerminalLifecycle lifecycle = new TerminalLifecycle(false, false);
        JdbcFederationResultCursor cursor = streamingCursor(
            lifecycle,
            cancellationGuard(cancelled)
        );

        InputStream stream = cursor.getBinaryStream(1);
        cancelled.set(true);
        FederationSqlException failure = expectStreamFailure(stream);

        Assert.assertEquals(FederationSqlErrorCode.QUERY_CANCELLED, failure.errorCode());
        Assert.assertTrue(lifecycle.unregistered.get());
    }

    /**
     * 验证字符流取得后发生超时时，后续流读取立即终止并释放 JDBC 资源。
     *
     * @throws Exception 流读取失败
     */
    @Test
    public void shouldStopCharacterStreamReadAfterTimeout() throws Exception {
        AtomicBoolean timedOut = new AtomicBoolean();
        TerminalLifecycle lifecycle = new TerminalLifecycle(false, false);
        JdbcFederationResultCursor cursor = streamingCursor(
            lifecycle,
            timeoutGuard(timedOut)
        );

        Reader reader = cursor.getCharacterStream(2);
        timedOut.set(true);
        FederationSqlException failure = expectStreamFailure(reader);

        Assert.assertEquals(FederationSqlErrorCode.QUERY_TIMEOUT, failure.errorCode());
        Assert.assertTrue(lifecycle.unregistered.get());
    }

    /**
     * 验证阻塞的结果读取可由 Statement.cancel 解阻，并确定性释放全部 JDBC 资源。
     *
     * @throws Exception 并发测试等待失败
     */
    @Test
    public void shouldCancelBlockingResultReadAndCloseAllResources() throws Exception {
        CountDownLatch readStarted = new CountDownLatch(1);
        CountDownLatch cancelSignal = new CountDownLatch(1);
        AtomicBoolean resultSetClosed = new AtomicBoolean();
        AtomicBoolean statementClosed = new AtomicBoolean();
        AtomicBoolean connectionClosed = new AtomicBoolean();
        CancellableLifecycle lifecycle = new CancellableLifecycle();
        DataSource dataSource = blockingReadDataSource(
            readStarted,
            cancelSignal,
            resultSetClosed,
            statementClosed,
            connectionClosed
        );
        JdbcFederationResultCursor cursor = (JdbcFederationResultCursor)
            new JdbcFederationFragmentExecutor().execute(context(
                FederationExecutionGuard.none(),
                dataSource,
                lifecycle
            ));
        ExecutorService executor = Executors.newSingleThreadExecutor();
        try {
            Future<FederationSqlException> read = executor.submit(
                () -> expectCursorFailure(cursor)
            );
            Assert.assertTrue(readStarted.await(2, TimeUnit.SECONDS));

            lifecycle.requestCancellation();
            FederationSqlException failure = read.get(2, TimeUnit.SECONDS);

            Assert.assertEquals(
                FederationSqlErrorCode.QUERY_CANCELLED,
                failure.errorCode()
            );
            Assert.assertTrue(lifecycle.unregistered.get());
            Assert.assertTrue(resultSetClosed.get());
            Assert.assertTrue(statementClosed.get());
            Assert.assertTrue(connectionClosed.get());
        } finally {
            executor.shutdownNow();
            executor.awaitTermination(2, TimeUnit.SECONDS);
        }
    }

    private static FederationSqlException expectFailure(
        FederationFragmentExecutionContext context
    ) {
        try {
            new JdbcFederationFragmentExecutor().execute(context);
            Assert.fail("expected connection acquisition to fail");
            return null;
        } catch (FederationSqlException exception) {
            return exception;
        }
    }

    private static FederationFragmentExecutionContext context(
        FederationExecutionGuard guard
    ) {
        return context(guard, new FailingDataSource(), new TerminalLifecycle(false, false));
    }

    private static FederationFragmentExecutionContext context(
        FederationExecutionGuard guard,
        DataSource dataSource,
        StatementLifecycle lifecycle
    ) {
        return new FederationFragmentExecutionContext(
            QueryId.create(),
            "SELECT 1",
            List.of(),
            SqlExecutionOptions.defaults(),
            dataSource,
            new AdapterCompatibility(
                AdapterCompatibilityStatus.VERIFIED,
                "test",
                "1",
                "test",
                "1",
                "test"
            ),
            Map.of(),
            lifecycle,
            null,
            guard
        );
    }

    private static DataSource executionTimeoutDataSource(
        AtomicBoolean statementClosed,
        AtomicBoolean connectionClosed
    ) {
        PreparedStatement statement = (PreparedStatement) Proxy.newProxyInstance(
            JdbcFederationFragmentExecutorTest.class.getClassLoader(),
            new Class<?>[] {PreparedStatement.class},
            (proxy, method, arguments) -> {
                if ("executeQuery".equals(method.getName())) {
                    throw new SQLTimeoutException("driver reported timeout after cancel");
                }
                if ("close".equals(method.getName())) {
                    statementClosed.set(true);
                }
                return defaultValue(method.getReturnType());
            }
        );
        Connection connection = (Connection) Proxy.newProxyInstance(
            JdbcFederationFragmentExecutorTest.class.getClassLoader(),
            new Class<?>[] {Connection.class},
            (proxy, method, arguments) -> {
                if ("prepareStatement".equals(method.getName())) {
                    return statement;
                }
                if ("close".equals(method.getName())) {
                    connectionClosed.set(true);
                }
                return defaultValue(method.getReturnType());
            }
        );
        return dataSource(connection);
    }

    /**
     * 创建在 ResultSet.next 中等待 Statement.cancel 的 JDBC 代理。
     *
     * @param readStarted 结果读取已开始信号
     * @param cancelSignal Statement 已取消信号
     * @param resultSetClosed ResultSet 关闭标记
     * @param statementClosed Statement 关闭标记
     * @param connectionClosed Connection 关闭标记
     * @return 可执行阻塞读取的 DataSource
     */
    private static DataSource blockingReadDataSource(
        CountDownLatch readStarted,
        CountDownLatch cancelSignal,
        AtomicBoolean resultSetClosed,
        AtomicBoolean statementClosed,
        AtomicBoolean connectionClosed
    ) {
        Object metadata = Proxy.newProxyInstance(
            JdbcFederationFragmentExecutorTest.class.getClassLoader(),
            new Class<?>[] {java.sql.ResultSetMetaData.class},
            (proxy, method, arguments) -> defaultValue(method.getReturnType())
        );
        ResultSet resultSet = (ResultSet) Proxy.newProxyInstance(
            JdbcFederationFragmentExecutorTest.class.getClassLoader(),
            new Class<?>[] {ResultSet.class},
            (proxy, method, arguments) -> {
                if ("getMetaData".equals(method.getName())) {
                    return metadata;
                }
                if ("next".equals(method.getName())) {
                    readStarted.countDown();
                    try {
                        if (!cancelSignal.await(2, TimeUnit.SECONDS)) {
                            throw new SQLException("test cancellation did not arrive");
                        }
                    } catch (InterruptedException exception) {
                        Thread.currentThread().interrupt();
                        throw new SQLException("blocking read was interrupted", exception);
                    }
                    throw new SQLException("driver read cancelled");
                }
                if ("close".equals(method.getName())) {
                    resultSetClosed.set(true);
                }
                return defaultValue(method.getReturnType());
            }
        );
        PreparedStatement statement = (PreparedStatement) Proxy.newProxyInstance(
            JdbcFederationFragmentExecutorTest.class.getClassLoader(),
            new Class<?>[] {PreparedStatement.class},
            (proxy, method, arguments) -> {
                if ("executeQuery".equals(method.getName())) {
                    return resultSet;
                }
                if ("cancel".equals(method.getName())) {
                    cancelSignal.countDown();
                }
                if ("close".equals(method.getName())) {
                    statementClosed.set(true);
                }
                return defaultValue(method.getReturnType());
            }
        );
        Connection connection = (Connection) Proxy.newProxyInstance(
            JdbcFederationFragmentExecutorTest.class.getClassLoader(),
            new Class<?>[] {Connection.class},
            (proxy, method, arguments) -> {
                if ("isReadOnly".equals(method.getName())) {
                    return true;
                }
                if ("prepareStatement".equals(method.getName())) {
                    return statement;
                }
                if ("close".equals(method.getName())) {
                    connectionClosed.set(true);
                }
                return defaultValue(method.getReturnType());
            }
        );
        return dataSource(connection);
    }

    private static JdbcFederationResultCursor failingCursor(
        StatementLifecycle lifecycle,
        SQLException readFailure
    ) {
        ResultSet resultSet = (ResultSet) Proxy.newProxyInstance(
            JdbcFederationFragmentExecutorTest.class.getClassLoader(),
            new Class<?>[] {ResultSet.class},
            (proxy, method, arguments) -> {
                if ("next".equals(method.getName())) {
                    throw readFailure;
                }
                return defaultValue(method.getReturnType());
            }
        );
        PreparedStatement statement = (PreparedStatement) Proxy.newProxyInstance(
            JdbcFederationFragmentExecutorTest.class.getClassLoader(),
            new Class<?>[] {PreparedStatement.class},
            (proxy, method, arguments) -> defaultValue(method.getReturnType())
        );
        Connection connection = (Connection) Proxy.newProxyInstance(
            JdbcFederationFragmentExecutorTest.class.getClassLoader(),
            new Class<?>[] {Connection.class},
            (proxy, method, arguments) -> defaultValue(method.getReturnType())
        );
        return new JdbcFederationResultCursor(
            QueryId.create(),
            List.of(),
            resultSet,
            statement,
            connection,
            lifecycle
        );
    }

    /**
     * 创建可返回二进制流和字符流的测试游标。
     *
     * @param lifecycle Statement 生命周期
     * @param guard 查询终态检查器
     * @return 测试游标
     */
    private static JdbcFederationResultCursor streamingCursor(
        StatementLifecycle lifecycle,
        FederationExecutionGuard guard
    ) {
        ResultSet resultSet = (ResultSet) Proxy.newProxyInstance(
            JdbcFederationFragmentExecutorTest.class.getClassLoader(),
            new Class<?>[] {ResultSet.class},
            (proxy, method, arguments) -> {
                if ("getBinaryStream".equals(method.getName())) {
                    return new ByteArrayInputStream(new byte[] {1, 2, 3});
                }
                if ("getCharacterStream".equals(method.getName())) {
                    return new StringReader("streamed value");
                }
                return defaultValue(method.getReturnType());
            }
        );
        PreparedStatement statement = (PreparedStatement) Proxy.newProxyInstance(
            JdbcFederationFragmentExecutorTest.class.getClassLoader(),
            new Class<?>[] {PreparedStatement.class},
            (proxy, method, arguments) -> defaultValue(method.getReturnType())
        );
        Connection connection = (Connection) Proxy.newProxyInstance(
            JdbcFederationFragmentExecutorTest.class.getClassLoader(),
            new Class<?>[] {Connection.class},
            (proxy, method, arguments) -> defaultValue(method.getReturnType())
        );
        return new JdbcFederationResultCursor(
            QueryId.create(),
            List.of(),
            resultSet,
            statement,
            connection,
            lifecycle,
            guard,
            FederationExecutionObserver.none()
        );
    }

    /**
     * 创建由布尔终态驱动的取消检查器。
     *
     * @param cancelled 是否已取消
     * @return 取消检查器
     */
    private static FederationExecutionGuard cancellationGuard(AtomicBoolean cancelled) {
        return terminalGuard(
            cancelled,
            FederationSqlErrorCode.QUERY_CANCELLED,
            "query was cancelled"
        );
    }

    /**
     * 创建由布尔终态驱动的超时检查器。
     *
     * @param timedOut 是否已超时
     * @return 超时检查器
     */
    private static FederationExecutionGuard timeoutGuard(AtomicBoolean timedOut) {
        return terminalGuard(
            timedOut,
            FederationSqlErrorCode.QUERY_TIMEOUT,
            "query deadline reached"
        );
    }

    /**
     * 创建固定错误语义的查询终态检查器。
     *
     * @param terminal 是否进入终态
     * @param errorCode 终态错误码
     * @param message 错误消息
     * @return 查询终态检查器
     */
    private static FederationExecutionGuard terminalGuard(
        AtomicBoolean terminal,
        FederationSqlErrorCode errorCode,
        String message
    ) {
        return new FederationExecutionGuard() {
            @Override
            public void ensureAllowed() {
                if (terminal.get()) {
                    throw new FederationSqlException(errorCode, message);
                }
            }

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

    /**
     * 读取二进制流并捕获预期的统一异常。
     *
     * @param stream 测试流
     * @return 捕获的统一异常
     * @throws Exception 非预期读取错误
     */
    private static FederationSqlException expectStreamFailure(InputStream stream)
        throws Exception {
        try {
            stream.read();
            Assert.fail("expected binary stream read to fail");
            return null;
        } catch (FederationSqlException exception) {
            return exception;
        }
    }

    /**
     * 读取字符流并捕获预期的统一异常。
     *
     * @param reader 测试 Reader
     * @return 捕获的统一异常
     * @throws Exception 非预期读取错误
     */
    private static FederationSqlException expectStreamFailure(Reader reader)
        throws Exception {
        try {
            reader.read();
            Assert.fail("expected character stream read to fail");
            return null;
        } catch (FederationSqlException exception) {
            return exception;
        }
    }

    private static FederationSqlException expectCursorFailure(
        JdbcFederationResultCursor cursor
    ) {
        try {
            cursor.next();
            Assert.fail("expected result read to fail");
            return null;
        } catch (FederationSqlException exception) {
            return exception;
        }
    }

    private static DataSource dataSource(Connection connection) {
        return (DataSource) Proxy.newProxyInstance(
            JdbcFederationFragmentExecutorTest.class.getClassLoader(),
            new Class<?>[] {DataSource.class},
            (proxy, method, arguments) -> {
                if ("getConnection".equals(method.getName())) {
                    return connection;
                }
                if ("getParentLogger".equals(method.getName())) {
                    return Logger.getGlobal();
                }
                return defaultValue(method.getReturnType());
            }
        );
    }

    private static Object defaultValue(Class<?> type) {
        if (!type.isPrimitive()) {
            return null;
        }
        if (type == boolean.class) {
            return false;
        }
        if (type == byte.class) {
            return (byte) 0;
        }
        if (type == short.class) {
            return (short) 0;
        }
        if (type == int.class) {
            return 0;
        }
        if (type == long.class) {
            return 0L;
        }
        if (type == float.class) {
            return 0F;
        }
        if (type == double.class) {
            return 0D;
        }
        if (type == char.class) {
            return '\0';
        }
        return null;
    }

    /** 记录测试所需的查询终态和注销动作。 */
    private static final class TerminalLifecycle implements StatementLifecycle {

        private final boolean timedOut;
        private final boolean cancelled;
        private final AtomicBoolean unregistered = new AtomicBoolean();

        /**
         * 创建固定终态的生命周期。
         *
         * @param timedOut 是否已超时
         * @param cancelled 是否已取消
         */
        private TerminalLifecycle(boolean timedOut, boolean cancelled) {
            this.timedOut = timedOut;
            this.cancelled = cancelled;
        }

        /** {@inheritDoc} */
        @Override
        public void register(Statement statement) {
        }

        /** {@inheritDoc} */
        @Override
        public void unregister(Statement statement) {
            unregistered.set(true);
        }

        /** {@inheritDoc} */
        @Override
        public boolean cancellationRequested() {
            return cancelled;
        }

        /** {@inheritDoc} */
        @Override
        public boolean timeoutRequested() {
            return timedOut;
        }
    }

    /** 可从测试线程触发 Statement.cancel 的生命周期。 */
    private static final class CancellableLifecycle implements StatementLifecycle {

        private final AtomicReference<Statement> statement = new AtomicReference<>();
        private final AtomicBoolean cancelled = new AtomicBoolean();
        private final AtomicBoolean unregistered = new AtomicBoolean();

        /** {@inheritDoc} */
        @Override
        public void register(Statement candidate) {
            statement.set(candidate);
        }

        /** {@inheritDoc} */
        @Override
        public void unregister(Statement candidate) {
            statement.compareAndSet(candidate, null);
            unregistered.set(true);
        }

        /** {@inheritDoc} */
        @Override
        public boolean cancellationRequested() {
            return cancelled.get();
        }

        /**
         * 标记查询取消并调用已登记 Statement 的取消入口。
         *
         * @throws SQLException JDBC 取消失败
         */
        private void requestCancellation() throws SQLException {
            cancelled.set(true);
            Statement active = statement.get();
            if (active != null) {
                active.cancel();
            }
        }
    }

    /** 始终返回瞬时连接池超时的测试 DataSource。 */
    private static final class FailingDataSource implements DataSource {

        @Override
        public Connection getConnection() throws SQLException {
            throw new SQLTransientConnectionException("pool timeout");
        }

        @Override
        public Connection getConnection(String username, String password) throws SQLException {
            throw new SQLTransientConnectionException("pool timeout");
        }

        @Override
        public PrintWriter getLogWriter() {
            return null;
        }

        @Override
        public void setLogWriter(PrintWriter out) {
        }

        @Override
        public void setLoginTimeout(int seconds) {
        }

        @Override
        public int getLoginTimeout() {
            return 0;
        }

        @Override
        public Logger getParentLogger() {
            return Logger.getGlobal();
        }

        @Override
        public <T> T unwrap(Class<T> iface) throws SQLException {
            throw new SQLException("unsupported");
        }

        @Override
        public boolean isWrapperFor(Class<?> iface) {
            return false;
        }
    }
}
