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

import com.easyagents.federation.sql.api.FederationSqlEngine;
import com.easyagents.federation.sql.api.FederationSqlEngines;
import com.easyagents.federation.sql.api.SqlQueryCommand;
import com.easyagents.federation.sql.execute.FederationResultCursor;
import com.easyagents.federation.sql.federation.FederationExecutionPolicy;
import com.easyagents.federation.sql.federation.FederationQueryScopeDefinition;
import com.easyagents.federation.sql.federation.FederationSourceBindingDefinition;
import com.easyagents.federation.sql.source.FederationDataSourceHandles;
import com.easyagents.federation.sql.source.FederationSourceDefinition;
import com.easyagents.federation.sql.source.RuntimeFingerprint;
import com.easyagents.federation.sql.source.SourceApplyOptions;
import com.easyagents.federation.sql.source.SourceId;
import com.mysql.cj.jdbc.MysqlDataSource;
import org.junit.Assert;
import org.junit.Assume;
import org.junit.Test;

import java.math.BigDecimal;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.List;
import java.util.Map;

/**
 * 亿级数据集 100,000 行 Cursor 手工基准；默认跳过，必须显式启用。
 */
public class Jdbc100mCursorManualBenchmarkTest {

    /**
     * 从支付库顺序消费恰好 100,000 行，验证 Cursor 流式消费和资源上限。
     */
    @Test
    public void shouldConsumeOneHundredThousandRows() {
        Assume.assumeTrue(Boolean.getBoolean("federation.100m.enabled"));
        String host = System.getProperty("federation.100m.mysql.host", "127.0.0.1");
        int port = Integer.getInteger("federation.100m.mysql.port", 33307);
        String database = System.getProperty(
            "federation.100m.mysql.database", "ef_bank_payment_perf");
        String username = System.getProperty("federation.100m.mysql.username", "ef_bench_ro");
        String password = System.getProperty("federation.100m.mysql.password");
        Assert.assertNotNull("必须通过系统属性提供只读密码", password);

        MysqlDataSource dataSource = new MysqlDataSource();
        dataSource.setUrl("jdbc:mysql://" + host + ':' + port + '/' + database
            + "?useSSL=false&allowPublicKeyRetrieval=true&serverTimezone=Asia/Shanghai"
            + "&useCursorFetch=true&defaultFetchSize=1000");
        dataSource.setUser(username);
        dataSource.setPassword(password);

        SourceId sourceId = new SourceId("payment-100m");
        FederationSourceDefinition definition = new FederationSourceDefinition(
            sourceId, 1, JdbcFederationSqlAdapterProvider.ADAPTER_ID,
            List.of(new JdbcSchemaDefinition("MAIN", database, null)), Map.of());
        FederationExecutionPolicy policy = new FederationExecutionPolicy(
            6, 12, 2, 100_000, 64L * 1024L * 1024L, 60_000);
        FederationQueryScopeDefinition scope = FederationQueryScopeDefinition.virtual(
            "payment-cursor-100m", 1,
            Map.of("PAY", FederationSourceBindingDefinition.of(sourceId, 1)),
            "PAY", policy);

        CursorResult nativeResult = consumeNative(dataSource);
        CursorResult federationResult;
        long federationStarted = System.nanoTime();
        try (FederationSqlEngine engine = FederationSqlEngines.builder()
            .dataSourceResolver(ignored -> FederationDataSourceHandles.shared(
                dataSource, new RuntimeFingerprint("MySQL", "8.0", "JDBC", "8.4", "1")))
            .adapter(new JdbcFederationSqlAdapterProvider())
            .federationExecutionPolicy(policy)
            .build()) {
            engine.sources().apply(definition, SourceApplyOptions.prewarmNow());
            String sql = "SELECT payment_id, amount FROM PAY.MAIN.payment_order "
                + "WHERE payment_id BETWEEN 1 AND 100000 ORDER BY payment_id";
            int rows = 0;
            long amountCents = 0;
            try (FederationResultCursor cursor = engine.query(
                SqlQueryCommand.of(sql, scope, List.of()))) {
                while (cursor.next()) {
                    rows++;
                    amountCents += ((BigDecimal) cursor.getObject(2))
                        .movePointRight(2).longValue();
                }
            }
            federationResult = new CursorResult(rows, amountCents,
                elapsedMillis(federationStarted));
        }
        Assert.assertEquals(100_000, nativeResult.rows());
        Assert.assertEquals(100_000, federationResult.rows());
        Assert.assertEquals(nativeResult.amountCents(), federationResult.amountCents());
        Assert.assertTrue(federationResult.amountCents() > 0L);
        System.out.println("BENCH_CURSOR_100K rows=" + federationResult.rows()
            + " nativeMillis=" + nativeResult.elapsedMillis()
            + " federationMillis=" + federationResult.elapsedMillis()
            + " amountCents=" + federationResult.amountCents());
    }

    private CursorResult consumeNative(MysqlDataSource dataSource) {
        String sql = "SELECT payment_id, amount FROM payment_order "
            + "WHERE payment_id BETWEEN 1 AND 100000 ORDER BY payment_id";
        long started = System.nanoTime();
        int rows = 0;
        long amountCents = 0;
        try (Connection connection = dataSource.getConnection();
             PreparedStatement statement = connection.prepareStatement(sql,
                 ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY)) {
            statement.setFetchSize(1000);
            try (ResultSet resultSet = statement.executeQuery()) {
                while (resultSet.next()) {
                    rows++;
                    amountCents += resultSet.getBigDecimal(2)
                        .movePointRight(2).longValue();
                }
            }
        } catch (SQLException exception) {
            throw new IllegalStateException("原生 JDBC Cursor 基准执行失败", exception);
        }
        return new CursorResult(rows, amountCents, elapsedMillis(started));
    }

    private long elapsedMillis(long started) {
        return (System.nanoTime() - started) / 1_000_000L;
    }

    private record CursorResult(int rows, long amountCents, long elapsedMillis) {
    }
}
