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.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import com.easyagents.federation.sql.api.SqlQueryCommand;
import com.easyagents.federation.sql.compile.FederationSqlPlan;
import com.easyagents.federation.sql.compile.SqlCompileRequest;
import com.easyagents.federation.sql.compile.SqlExplainLevel;
import com.easyagents.federation.sql.compile.SqlExplainRequest;
import com.easyagents.federation.sql.compile.SqlExplainResult;
import com.easyagents.federation.sql.execute.FederationQueryMetricsSnapshot;
import com.easyagents.federation.sql.execute.FederationResultCursor;
import com.easyagents.federation.sql.execute.SqlParameter;
import com.easyagents.federation.sql.federation.FederationExecutionPolicy;
import com.easyagents.federation.sql.federation.FederationLogicalTableDefinition;
import com.easyagents.federation.sql.federation.FederationQueryMode;
import com.easyagents.federation.sql.federation.FederationQueryScopeDefinition;
import com.easyagents.federation.sql.federation.FederationSourceBindingDefinition;
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.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 java.math.BigDecimal;
import java.time.Instant;
import java.sql.Connection;
import java.sql.Statement;
import java.sql.Types;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import org.h2.jdbcx.JdbcDataSource;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

/**
 * 多个独立 JDBC DataSource 的基础联邦查询集成测试。
 */
public class JdbcFederatedQueryEngineTest {

    private static final SourceId SALES_SOURCE = new SourceId("sales-source");
    private static final SourceId BILLING_SOURCE = new SourceId("billing-source");
    private static final SourceId REGION_SOURCE = new SourceId("region-source");

    private JdbcDataSource sales;
    private JdbcDataSource billing;
    private JdbcDataSource region;
    private FederationSqlEngine engine;
    private FederationQueryScopeDefinition scope;
    private final AtomicReference<String> statisticsVersion =
        new AtomicReference<>("stats-v1");
    private final AtomicLong salesRowCount = new AtomicLong(3);

    /**
     * 创建三个独立 H2 数据库并登记物理数据源。
     *
     * @throws Exception 数据库初始化失败
     */
    @Before
    public void setUp() throws Exception {
        Instant statisticsCollectedAt = Instant.now();
        sales = dataSource("sales");
        billing = dataSource("billing");
        region = dataSource("region");
        execute(sales,
            "CREATE TABLE CUSTOMER (ID INT PRIMARY KEY, NAME VARCHAR(64) NOT NULL)",
            "INSERT INTO CUSTOMER VALUES (1, 'Alice'), (2, 'Bob'), (3, 'Carol')",
            "CREATE TABLE TIME_EVENT (ID INT PRIMARY KEY, EVENT_TIME TIME(6) WITH TIME ZONE, "
                + "EVENT_AT TIMESTAMP(6) WITH TIME ZONE)",
            "INSERT INTO TIME_EVENT VALUES (1, TIME WITH TIME ZONE '12:00:00.123456+08:00', "
                + "TIMESTAMP WITH TIME ZONE '2026-08-21 12:00:00.123456+08:00')",
            "CREATE TABLE PRECISE_EVENT (ID INT PRIMARY KEY, EVENT_AT TIMESTAMP(6))",
            "INSERT INTO PRECISE_EVENT VALUES "
                + "(1, TIMESTAMP '2026-08-21 12:00:00.123456')");
        execute(billing,
            "CREATE TABLE ORDER_ITEM (ID INT PRIMARY KEY, CUSTOMER_ID INT NOT NULL, AMOUNT DECIMAL(12,2), CODE VARCHAR(64))",
            "INSERT INTO ORDER_ITEM VALUES (10, 1, 30.00, 'Alice'), (11, 1, 20.00, 'Alice'), (12, 2, 80.00, 'Bob')",
            "CREATE TABLE TIME_EVENT (ID INT PRIMARY KEY, EVENT_TIME TIME(6) WITH TIME ZONE, "
                + "EVENT_AT TIMESTAMP(6) WITH TIME ZONE)",
            "INSERT INTO TIME_EVENT VALUES (2, TIME WITH TIME ZONE '06:00:00.123456+02:00', "
                + "TIMESTAMP WITH TIME ZONE '2026-08-21 06:00:00.123456+02:00')",
            "CREATE TABLE PRECISE_EVENT (ID INT PRIMARY KEY, EVENT_AT TIMESTAMP(6))",
            "INSERT INTO PRECISE_EVENT VALUES "
                + "(2, TIMESTAMP '2026-08-21 08:30:00.654321')");
        execute(region,
            "CREATE TABLE CUSTOMER_REGION (CUSTOMER_ID INT PRIMARY KEY, REGION VARCHAR(64))",
            "INSERT INTO CUSTOMER_REGION VALUES (1, 'North'), (2, 'South'), (3, 'West')");

        RuntimeFingerprint fingerprint = new RuntimeFingerprint(
            "H2", "2", "H2 JDBC Driver", "2", "1"
        );
        engine = FederationSqlEngines.builder()
            .dataSourceResolver(definition -> FederationDataSourceHandles.shared(
                dataSourceFor(definition.sourceId()),
                fingerprint
            ))
            .tableStatisticsProvider(() -> new FederationStatisticsSnapshot(
                statisticsVersion.get(),
                Map.of(
                    new FederationStatisticsSnapshot.TableKey(
                        SALES_SOURCE, "APP", "CUSTOMER"
                    ),
                    new FederationTableStatistics(
                        salesRowCount.get(),
                        48,
                        statisticsCollectedAt,
                        "test-catalog"
                    ),
                    new FederationStatisticsSnapshot.TableKey(
                        BILLING_SOURCE, "APP", "ORDER_ITEM"
                    ),
                    new FederationTableStatistics(
                        3,
                        64,
                        statisticsCollectedAt,
                        "test-catalog"
                    ),
                    new FederationStatisticsSnapshot.TableKey(
                        REGION_SOURCE, "APP", "CUSTOMER_REGION"
                    ),
                    new FederationTableStatistics(
                        3,
                        32,
                        statisticsCollectedAt,
                        "test-catalog"
                    )
                )
            ))
            .federationExecutionPolicy(threeSourcePolicy())
            .maximumPlanCacheEntries(32)
            .build();
        engine.sources().apply(definition(SALES_SOURCE), SourceApplyOptions.prewarmNow());
        engine.sources().apply(definition(BILLING_SOURCE), SourceApplyOptions.prewarmNow());
        engine.sources().apply(definition(REGION_SOURCE), SourceApplyOptions.prewarmNow());

        scope = new FederationQueryScopeDefinition(
            "sales-billing",
            1,
            Map.of(
                "SALES",
                FederationSourceBindingDefinition.of(SALES_SOURCE, 1),
                "BILLING",
                FederationSourceBindingDefinition.of(BILLING_SOURCE, 1)
            ),
            "SALES",
            FederationExecutionPolicy.basic()
        );
    }

    /**
     * 关闭 Engine。
     */
    @After
    public void tearDown() {
        if (engine != null) {
            engine.close();
        }
    }

    /**
     * 验证同一 Query Scope 中实际只引用一个源时保持完整单源下推。
     */
    @Test
    public void shouldRouteSingleReferencedSourceToDirectExecution() {
        FederationSqlPlan plan = engine.compile(SqlCompileRequest.of(
            "SELECT NAME FROM SALES.APP.CUSTOMER ORDER BY ID",
            scope
        ));

        Assert.assertEquals(FederationQueryMode.SINGLE_SOURCE, plan.queryMode());
        Assert.assertEquals(1, plan.fragments().size());
        Assert.assertEquals(java.util.Set.of(SALES_SOURCE), plan.referencedSources());
    }

    /**
     * 验证短逻辑表名、三段逻辑表名和跨源逻辑表 Join 共用底层映射。
     */
    @Test
    public void shouldExecuteLogicalTableNamesAcrossSources() {
        FederationQueryScopeDefinition logicalScope =
            FederationQueryScopeDefinition.virtual(
                "logical-sales-billing",
                2,
                scope.bindings(),
                "SALES",
                List.of(
                    FederationLogicalTableDefinition.of(
                        "customers", "SALES", "APP", "CUSTOMER"
                    ),
                    FederationLogicalTableDefinition.of(
                        "order_lines", "BILLING", "APP", "ORDER_ITEM"
                    )
                ),
                FederationExecutionPolicy.basic()
            );

        try (FederationResultCursor cursor = engine.query(SqlQueryCommand.of(
            "SELECT customers.NAME FROM customers ORDER BY customers.ID",
            logicalScope,
            List.of()
        ))) {
            Assert.assertTrue(cursor.next());
            Assert.assertEquals("Alice", cursor.getObject(1));
        }

        try (FederationResultCursor cursor = engine.query(SqlQueryCommand.of(
            "SELECT customers.NAME FROM SALES.APP.customers "
                + "ORDER BY SALES.APP.customers.ID",
            logicalScope,
            List.of()
        ))) {
            Assert.assertTrue(cursor.next());
            Assert.assertEquals("Alice", cursor.getObject(1));
        }

        try (FederationResultCursor cursor = engine.query(SqlQueryCommand.of(
            "SELECT c.NAME, o.AMOUNT FROM customers c "
                + "JOIN order_lines o ON c.ID = o.CUSTOMER_ID "
                + "ORDER BY o.ID",
            logicalScope,
            List.of()
        ))) {
            Assert.assertTrue(cursor.next());
            Assert.assertEquals(List.of("Alice", new BigDecimal("30.00")), cursor.row());
        }
    }

    /**
     * 验证全限定 SQL 不依赖未引用默认 Binding 的运行状态。
     */
    @Test
    public void shouldNotAcquireUnavailableDefaultBindingWhenOnlyAnotherSourceIsReferenced() {
        FederationQueryScopeDefinition unavailableDefault =
            FederationQueryScopeDefinition.virtual(
                "unavailable-default",
                1,
                Map.of(
                    "OFFLINE", FederationSourceBindingDefinition.of(
                        new SourceId("offline-source"), 1
                    ),
                    "BILLING", FederationSourceBindingDefinition.of(BILLING_SOURCE, 1)
                ),
                "OFFLINE",
                FederationExecutionPolicy.basic()
            );

        FederationSqlPlan plan = engine.compile(SqlCompileRequest.of(
            "SELECT AMOUNT FROM BILLING.APP.ORDER_ITEM",
            unavailableDefault
        ));

        Assert.assertEquals(FederationQueryMode.SINGLE_SOURCE, plan.queryMode());
        Assert.assertEquals(java.util.Set.of(BILLING_SOURCE), plan.referencedSources());

        FederationSqlPlan ctePlan = engine.compile(SqlCompileRequest.of(
            "WITH x AS (SELECT AMOUNT FROM BILLING.APP.ORDER_ITEM) SELECT * FROM x",
            unavailableDefault
        ));
        Assert.assertEquals(FederationQueryMode.SINGLE_SOURCE, ctePlan.queryMode());
        Assert.assertEquals(java.util.Set.of(BILLING_SOURCE), ctePlan.referencedSources());
    }

    /**
     * 验证跨源等值 Join、聚合、全局排序和查询指标。
     */
    @Test
    public void shouldJoinAggregateAndSortAcrossTwoSources() {
        String sql = "SELECT c.ID, SUM(o.AMOUNT) AS TOTAL "
            + "FROM SALES.APP.CUSTOMER c "
            + "JOIN BILLING.APP.ORDER_ITEM o ON c.ID = o.CUSTOMER_ID "
            + "GROUP BY c.ID ORDER BY TOTAL DESC";
        FederationSqlPlan plan = engine.compile(SqlCompileRequest.of(sql, scope));
        Assert.assertEquals(FederationQueryMode.FEDERATED, plan.queryMode());
        Assert.assertEquals(2, plan.fragments().size());
        Assert.assertEquals(
            "the smaller estimated input should become the Hash Join build side",
            SALES_SOURCE,
            plan.fragments().get(1).sourceId()
        );

        List<List<Object>> rows = new ArrayList<>();
        FederationQueryMetricsSnapshot finalMetrics;
        try (FederationResultCursor cursor = engine.query(
            SqlQueryCommand.of(sql, scope, List.of())
        )) {
            while (cursor.next()) {
                rows.add(cursor.row());
            }
            finalMetrics = cursor.metrics();
        }

        Assert.assertEquals(2, rows.size());
        Assert.assertEquals(2, rows.get(0).get(0));
        Assert.assertEquals(new BigDecimal("80.00"), rows.get(0).get(1));
        Assert.assertEquals(1, rows.get(1).get(0));
        Assert.assertEquals(new BigDecimal("50.00"), rows.get(1).get(1));
        Assert.assertTrue(finalMetrics.complete());
        Assert.assertEquals(2, finalMetrics.returnedRows());
        Assert.assertTrue(finalMetrics.intermediateRows() >= 6);
        Assert.assertEquals(2, finalMetrics.fragments().size());
    }

    /**
     * 验证三个独立 JDBC 数据源可由同一计划完成 Join 并返回稳定结果。
     */
    @Test
    public void shouldExecuteJoinAcrossThreeSources() {
        FederationQueryScopeDefinition threeSourceScope =
            new FederationQueryScopeDefinition(
                "sales-billing-region",
                1,
                Map.of(
                    "SALES", FederationSourceBindingDefinition.of(SALES_SOURCE, 1),
                    "BILLING", FederationSourceBindingDefinition.of(BILLING_SOURCE, 1),
                    "REGION", FederationSourceBindingDefinition.of(REGION_SOURCE, 1)
                ),
                "SALES",
                threeSourcePolicy()
            );
        String sql = "SELECT c.NAME, o.AMOUNT, r.REGION "
            + "FROM SALES.APP.CUSTOMER c "
            + "JOIN BILLING.APP.ORDER_ITEM o ON c.ID = o.CUSTOMER_ID "
            + "JOIN REGION.APP.CUSTOMER_REGION r ON c.ID = r.CUSTOMER_ID "
            + "ORDER BY o.ID";

        FederationSqlPlan plan = engine.compile(SqlCompileRequest.of(sql, threeSourceScope));
        Assert.assertEquals(FederationQueryMode.FEDERATED, plan.queryMode());
        Assert.assertEquals(3, plan.referencedSources().size());
        Assert.assertEquals(3, plan.fragments().size());

        List<List<Object>> rows = new ArrayList<>();
        try (FederationResultCursor cursor = engine.query(
            SqlQueryCommand.of(sql, threeSourceScope, List.of())
        )) {
            while (cursor.next()) {
                rows.add(cursor.row());
            }
        }
        Assert.assertEquals(3, rows.size());
        Assert.assertEquals(
            List.of("Alice", new BigDecimal("30.00"), "North"),
            rows.get(0)
        );
        Assert.assertEquals(
            List.of("Bob", new BigDecimal("80.00"), "South"),
            rows.get(2)
        );
    }

    /**
     * 验证两个分片分别绑定原始查询中的动态参数。
     */
    @Test
    public void shouldMapParametersIntoDifferentFragments() {
        String sql = "SELECT c.NAME, o.AMOUNT "
            + "FROM SALES.APP.CUSTOMER c "
            + "JOIN BILLING.APP.ORDER_ITEM o ON c.ID = o.CUSTOMER_ID "
            + "WHERE c.ID > ? AND o.AMOUNT > ? ORDER BY o.AMOUNT";
        FederationSqlPlan plan = engine.compile(new SqlCompileRequest(
            sql,
            scope,
            List.of(Types.INTEGER, Types.DECIMAL),
            "default"
        ));
        Assert.assertEquals(2, plan.fragments().size());
        Assert.assertEquals(List.of(0), plan.fragments().get(0).parameterMapping());
        Assert.assertEquals(List.of(1), plan.fragments().get(1).parameterMapping());
        Assert.assertTrue(plan.fragments().stream()
            .allMatch(fragment -> fragment.executableSql().toUpperCase().contains("WHERE")));
        try (FederationResultCursor cursor = engine.query(SqlQueryCommand.of(
            sql,
            scope,
            List.of(
                new SqlParameter(Types.INTEGER, 0),
                new SqlParameter(Types.DECIMAL, new BigDecimal("25.00"))
            )
        ))) {
            Assert.assertTrue(cursor.next());
            Assert.assertEquals(List.of("Alice", new BigDecimal("30.00")), cursor.row());
            Assert.assertTrue(cursor.next());
            Assert.assertEquals(List.of("Bob", new BigDecimal("80.00")), cursor.row());
            Assert.assertFalse(cursor.next());
        }

        try (FederationResultCursor cursor = engine.query(SqlQueryCommand.of(
            "SELECT NAME FROM SALES.APP.CUSTOMER ORDER BY ID "
                + "OFFSET ? ROWS FETCH NEXT ? ROWS ONLY",
            scope,
            List.of(
                new SqlParameter(Types.INTEGER, 1),
                new SqlParameter(Types.INTEGER, 1)
            )
        ))) {
            Assert.assertTrue(cursor.next());
            Assert.assertEquals("Bob", cursor.getObject(1));
            Assert.assertFalse(cursor.next());
        }
    }

    /**
     * 验证 LEFT JOIN、UNION ALL、CTE 和全局分页使用同一联邦执行入口。
     */
    @Test
    public void shouldExecuteBasicFederatedOperators() {
        List<List<Object>> leftRows = query(
            "SELECT c.NAME, o.AMOUNT FROM SALES.APP.CUSTOMER c "
                + "LEFT JOIN BILLING.APP.ORDER_ITEM o ON c.ID = o.CUSTOMER_ID "
                + "ORDER BY c.ID, o.ID"
        );
        Assert.assertEquals(4, leftRows.size());
        Assert.assertEquals("Carol", leftRows.get(3).get(0));
        Assert.assertNull(leftRows.get(3).get(1));

        List<List<Object>> unionRows = query(
            "SELECT ID FROM SALES.APP.CUSTOMER "
                + "UNION ALL SELECT CUSTOMER_ID FROM BILLING.APP.ORDER_ITEM ORDER BY ID"
        );
        Assert.assertEquals(List.of(1, 1, 1, 2, 2, 3),
            unionRows.stream().map(row -> row.get(0)).toList());

        List<List<Object>> cteRows = query(
            "WITH large_orders AS ("
                + "SELECT CUSTOMER_ID, AMOUNT FROM BILLING.APP.ORDER_ITEM WHERE AMOUNT >= 50"
                + ") SELECT c.NAME, o.AMOUNT FROM SALES.APP.CUSTOMER c "
                + "JOIN large_orders o ON c.ID = o.CUSTOMER_ID "
                + "ORDER BY o.AMOUNT DESC FETCH NEXT 1 ROWS ONLY"
        );
        Assert.assertEquals(List.of(List.of("Bob", new BigDecimal("80.00"))), cteRows);

        List<List<Object>> residualRows = query(
            "SELECT c.ID, c.NAME FROM SALES.APP.CUSTOMER c "
                + "JOIN BILLING.APP.ORDER_ITEM o ON c.ID = o.CUSTOMER_ID "
                + "WHERE c.ID < o.ID ORDER BY c.ID, o.ID"
        );
        Assert.assertEquals(List.of("Alice", "Alice", "Bob"),
            residualRows.stream().map(row -> row.get(1)).toList());
    }

    /**
     * 验证本地联邦算子以 UTC Offset 类型返回 JDBC 4.2 时区值。
     */
    @Test
    public void shouldPreserveTimezoneSemanticsAcrossFederatedUnion() {
        List<List<Object>> rows = query(
            "SELECT ID, EVENT_TIME, EVENT_AT FROM SALES.APP.TIME_EVENT "
                + "UNION ALL SELECT ID, EVENT_TIME, EVENT_AT FROM BILLING.APP.TIME_EVENT "
                + "ORDER BY ID"
        );

        Assert.assertEquals(2, rows.size());
        Assert.assertEquals(
            java.time.OffsetTime.parse("04:00:00.123456Z"),
            rows.get(0).get(1)
        );
        Assert.assertEquals(
            java.time.OffsetDateTime.parse("2026-08-21T04:00:00.123456Z"),
            rows.get(0).get(2)
        );
        Assert.assertEquals(
            java.time.OffsetTime.parse("04:00:00.123456Z"),
            rows.get(1).get(1)
        );
        Assert.assertEquals(
            java.time.OffsetDateTime.parse("2026-08-21T04:00:00.123456Z"),
            rows.get(1).get(2)
        );

        List<List<Object>> joined = query(
            "SELECT s.ID, b.ID FROM SALES.APP.TIME_EVENT s "
                + "JOIN BILLING.APP.TIME_EVENT b ON s.EVENT_AT = b.EVENT_AT"
        );
        Assert.assertEquals(List.of(List.of(1, 2)), joined);
    }

    /**
     * 验证逻辑 Explain 不访问数据库 Optimizer，物理 Explain 显式返回每个分片原生计划。
     */
    @Test
    public void shouldExplainLogicalAndPhysicalFederatedPlans() {
        String sql = "SELECT c.NAME, o.AMOUNT FROM SALES.APP.CUSTOMER c "
            + "JOIN BILLING.APP.ORDER_ITEM o ON c.ID = o.CUSTOMER_ID";
        SqlCompileRequest compileRequest = SqlCompileRequest.of(sql, scope);

        SqlExplainResult logical = engine.explain(new SqlExplainRequest(
            compileRequest,
            SqlExplainLevel.LOGICAL
        ));
        Assert.assertEquals(FederationQueryMode.FEDERATED, logical.queryMode());
        Assert.assertEquals(2, logical.fragments().size());
        Assert.assertTrue(logical.fragments().stream()
            .allMatch(fragment -> fragment.physicalExplain() == null));
        Assert.assertTrue(logical.fragments().stream().allMatch(fragment ->
            fragment.costEstimate().estimatedRows() >= 0
                && fragment.costEstimate().estimatedRowWidthBytes() > 0
                && fragment.costEstimate().estimatedTransferBytes() >= 0
                && fragment.costEstimate().statisticsSource().contains("test-catalog")
                && "stats-v1".equals(
                    fragment.costEstimate().statisticsSnapshotVersion()
                )
                && !fragment.pushedDownOperators().isEmpty()
        ));

        SqlExplainResult physical = engine.explain(new SqlExplainRequest(compileRequest));
        Assert.assertEquals(SqlExplainLevel.PHYSICAL, physical.level());
        Assert.assertEquals(2, physical.fragments().size());
        Assert.assertTrue(physical.fragments().stream().allMatch(fragment ->
            fragment.physicalExplain() != null
                && fragment.physicalExplain().available()
                && !fragment.physicalExplain().nativePlan().isBlank()
        ));
    }

    /**
     * 验证纯快照版本变化不扰动计划，实际引用表统计变化才触发重编译。
     */
    @Test
    public void shouldRecompileWhenStatisticsSnapshotChanges() {
        SqlCompileRequest request = SqlCompileRequest.of(
            "SELECT NAME FROM SALES.APP.CUSTOMER ORDER BY ID",
            scope
        );
        FederationSqlPlan first = engine.compile(request);
        FederationSqlPlan cacheHit = engine.compile(request);
        Assert.assertSame(first.relRoot(), cacheHit.relRoot());

        statisticsVersion.set("stats-v2");
        FederationSqlPlan versionOnly = engine.compile(request);
        Assert.assertSame(first.relRoot(), versionOnly.relRoot());

        salesRowCount.set(4);
        FederationSqlPlan refreshed = engine.compile(request);

        Assert.assertNotSame(first.relRoot(), refreshed.relRoot());
        Assert.assertEquals(
            "stats-v2",
            refreshed.fragments().get(0).costEstimate().statisticsSnapshotVersion()
        );
    }

    /**
     * 验证不支持的跨源算子和中间结果预算超限均返回稳定错误。
     */
    @Test
    public void shouldRejectUnsupportedOperatorsAndExceededBudget() {
        assertCompileError(
            "SELECT c.NAME FROM SALES.APP.CUSTOMER c "
                + "JOIN BILLING.APP.ORDER_ITEM o ON c.ID < o.CUSTOMER_ID",
            FederationSqlErrorCode.FEDERATION_OPERATOR_UNSUPPORTED
        );
        assertCompileError(
            "SELECT ID FROM SALES.APP.CUSTOMER "
                + "UNION SELECT CUSTOMER_ID FROM BILLING.APP.ORDER_ITEM",
            FederationSqlErrorCode.FEDERATION_OPERATOR_UNSUPPORTED
        );
        assertCompileError(
            "SELECT c.ID FROM SALES.APP.CUSTOMER c "
                + "JOIN BILLING.APP.ORDER_ITEM o ON c.NAME = o.CODE",
            FederationSqlErrorCode.FEDERATION_OPERATOR_UNSUPPORTED
        );
        assertCompileError(
            "SELECT c.ID FROM SALES.APP.CUSTOMER c "
                + "JOIN BILLING.APP.ORDER_ITEM o ON c.ID = o.CUSTOMER_ID "
                + "AND c.NAME < o.CODE",
            FederationSqlErrorCode.FEDERATION_OPERATOR_UNSUPPORTED
        );
        assertCompileError(
            "SELECT c.ID FROM SALES.APP.CUSTOMER c "
                + "JOIN BILLING.APP.ORDER_ITEM o ON c.ID = o.CUSTOMER_ID "
                + "AND c.ID < o.ID",
            FederationSqlErrorCode.FEDERATION_OPERATOR_UNSUPPORTED
        );
        assertCompileError(
            "SELECT c.ID FROM SALES.APP.CUSTOMER c "
                + "JOIN BILLING.APP.ORDER_ITEM o ON c.ID = o.CUSTOMER_ID "
                + "AND c.NAME LIKE o.CODE",
            FederationSqlErrorCode.FEDERATION_OPERATOR_UNSUPPORTED
        );
        assertCompileError(
            "SELECT c.NAME FROM SALES.APP.CUSTOMER c "
                + "JOIN BILLING.APP.ORDER_ITEM o ON c.ID = o.CUSTOMER_ID "
                + "ORDER BY c.NAME",
            FederationSqlErrorCode.FEDERATION_OPERATOR_UNSUPPORTED
        );
        assertCompileError(
            "SELECT c.NAME, COUNT(*) FROM SALES.APP.CUSTOMER c "
                + "JOIN BILLING.APP.ORDER_ITEM o ON c.ID = o.CUSTOMER_ID "
                + "GROUP BY c.NAME",
            FederationSqlErrorCode.FEDERATION_OPERATOR_UNSUPPORTED
        );
        assertQueryError(
            "SELECT EVENT_AT FROM SALES.APP.PRECISE_EVENT "
                + "UNION ALL SELECT EVENT_AT FROM BILLING.APP.PRECISE_EVENT",
            FederationSqlErrorCode.FEDERATION_OPERATOR_UNSUPPORTED
        );

        FederationQueryScopeDefinition strictScope = FederationQueryScopeDefinition.virtual(
            "strict-budget",
            2,
            scope.bindings(),
            scope.defaultBinding(),
            new FederationExecutionPolicy(2, 8, 2, 1, 1024, 60_000)
        );
        try (FederationResultCursor cursor = engine.query(SqlQueryCommand.of(
            "SELECT c.NAME, o.AMOUNT FROM SALES.APP.CUSTOMER c "
                + "JOIN BILLING.APP.ORDER_ITEM o ON c.ID = o.CUSTOMER_ID",
            strictScope,
            List.of()
        ))) {
            cursor.next();
            Assert.fail("intermediate row budget should reject the query");
        } catch (FederationSqlException exception) {
            Assert.assertEquals(
                FederationSqlErrorCode.FEDERATION_RESOURCE_LIMIT_EXCEEDED,
                exception.errorCode()
            );
        }

        FederationQueryScopeDefinition localExpansionScope =
            FederationQueryScopeDefinition.virtual(
                "local-expansion-budget",
                3,
                scope.bindings(),
                scope.defaultBinding(),
                new FederationExecutionPolicy(2, 8, 2, 7, 64_000, 60_000)
            );
        try (FederationResultCursor cursor = engine.query(SqlQueryCommand.of(
            "SELECT COUNT(*) FROM SALES.APP.CUSTOMER c "
                + "JOIN BILLING.APP.ORDER_ITEM o ON c.ID = o.CUSTOMER_ID",
            localExpansionScope,
            List.of()
        ))) {
            cursor.next();
            Assert.fail("local join expansion should consume the intermediate row budget");
        } catch (FederationSqlException exception) {
            Assert.assertEquals(
                FederationSqlErrorCode.FEDERATION_RESOURCE_LIMIT_EXCEEDED,
                exception.errorCode()
            );
        }

        Assert.assertFalse(query(
            "SELECT c.NAME FROM SALES.APP.CUSTOMER c "
                + "JOIN BILLING.APP.ORDER_ITEM o ON c.ID = o.CUSTOMER_ID"
        ).isEmpty());

        FederationQueryScopeDefinition singleFragmentSlot =
            FederationQueryScopeDefinition.virtual(
                "single-fragment-slot",
                3,
                scope.bindings(),
                scope.defaultBinding(),
                new FederationExecutionPolicy(2, 8, 1, 100, 64_000, 2_000)
            );
        try (FederationResultCursor cursor = engine.query(SqlQueryCommand.of(
            "SELECT c.NAME FROM SALES.APP.CUSTOMER c "
                + "JOIN BILLING.APP.ORDER_ITEM o ON c.ID = o.CUSTOMER_ID",
            singleFragmentSlot,
            List.of()
        ))) {
            Assert.assertTrue(cursor.next());
        }
    }

    private List<List<Object>> query(String sql) {
        List<List<Object>> rows = new ArrayList<>();
        try (FederationResultCursor cursor = engine.query(
            SqlQueryCommand.of(sql, scope, List.of())
        )) {
            while (cursor.next()) {
                rows.add(cursor.row());
            }
        }
        return rows;
    }

    private void assertCompileError(String sql, FederationSqlErrorCode expected) {
        try {
            engine.compile(SqlCompileRequest.of(sql, scope));
            Assert.fail("SQL should have been rejected: " + sql);
        } catch (FederationSqlException exception) {
            Assert.assertEquals(expected, exception.errorCode());
        }
    }

    /**
     * 断言联邦查询在执行阶段返回指定稳定错误。
     *
     * @param sql 待执行 SQL
     * @param expected 预期错误码
     */
    private void assertQueryError(String sql, FederationSqlErrorCode expected) {
        try (FederationResultCursor cursor = engine.query(
            SqlQueryCommand.of(sql, scope, List.of())
        )) {
            cursor.next();
            Assert.fail("SQL execution should have been rejected: " + sql);
        } catch (FederationSqlException exception) {
            Assert.assertEquals(expected, exception.errorCode());
        }
    }

    private static JdbcDataSource dataSource(String name) {
        JdbcDataSource dataSource = new JdbcDataSource();
        dataSource.setURL(
            "jdbc:h2:mem:federation_" + name + '_' + System.nanoTime()
                + ";DB_CLOSE_DELAY=-1"
        );
        return dataSource;
    }

    /**
     * 根据物理源选择测试数据库。
     *
     * @param sourceId 物理数据源标识
     * @return 对应测试数据源
     */
    private JdbcDataSource dataSourceFor(SourceId sourceId) {
        if (sourceId.equals(SALES_SOURCE)) {
            return sales;
        }
        if (sourceId.equals(BILLING_SOURCE)) {
            return billing;
        }
        if (sourceId.equals(REGION_SOURCE)) {
            return region;
        }
        throw new IllegalArgumentException("unknown test source: " + sourceId.value());
    }

    /**
     * 返回允许三源执行且限制并发分片数的测试策略。
     *
     * @return 三源执行策略
     */
    private static FederationExecutionPolicy threeSourcePolicy() {
        return new FederationExecutionPolicy(
            3,
            8,
            2,
            100_000,
            64L * 1024L * 1024L,
            60_000
        );
    }

    private static void execute(JdbcDataSource dataSource, String... statements)
        throws Exception {
        try (Connection connection = dataSource.getConnection();
             Statement statement = connection.createStatement()) {
            for (String sql : statements) {
                statement.execute(sql);
            }
        }
    }

    private static FederationSourceDefinition definition(SourceId sourceId) {
        return new FederationSourceDefinition(
            sourceId,
            1,
            JdbcFederationSqlAdapterProvider.ADAPTER_ID,
            List.of(new JdbcSchemaDefinition("APP", null, "PUBLIC")),
            Map.of()
        );
    }
}
