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.FederationFragmentExplainContext;
import java.io.PrintWriter;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.logging.Logger;
import javax.sql.DataSource;
import org.junit.Assert;
import org.junit.Test;

/**
 * JDBC 物理 Explain 的连接获取错误边界测试。
 */
public class JdbcFederationFragmentExplainerTest {

    /**
     * 验证连接池以运行时异常拒绝连接时返回精确连接获取失败错误。
     */
    @Test
    public void shouldMapRuntimeConnectionFailure() {
        FederationSqlException failure = expectFailure(context(
            FederationExecutionGuard.none(),
            new IllegalStateException("pool is closed")
        ));

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

    /**
     * 验证连接池失败返回时，已经到达的统一截止时间优先于连接错误。
     */
    @Test
    public void shouldPreferDeadlineOverRuntimeConnectionFailure() {
        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,
            new IllegalStateException("pool is closed")
        ));

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

    private static FederationSqlException expectFailure(
        FederationFragmentExplainContext context
    ) {
        try {
            new JdbcFederationFragmentExplainer().explain(context);
            Assert.fail("physical Explain should fail");
            return null;
        } catch (FederationSqlException exception) {
            return exception;
        }
    }

    private static FederationFragmentExplainContext context(
        FederationExecutionGuard guard,
        RuntimeException failure
    ) {
        return new FederationFragmentExplainContext(
            "SELECT 1",
            List.of(),
            failingDataSource(failure),
            new AdapterCompatibility(
                AdapterCompatibilityStatus.VERIFIED,
                "MySQL",
                "8",
                "test-driver",
                "1",
                "test"
            ),
            Map.of(),
            5,
            guard
        );
    }

    private static DataSource failingDataSource(RuntimeException failure) {
        return new DataSource() {
            @Override
            public Connection getConnection() {
                throw failure;
            }

            @Override
            public Connection getConnection(String username, String password) {
                throw failure;
            }

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

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

            @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();
            }
        };
    }
}
