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.execute.FederationResultCursor;
import com.easyagents.federation.sql.federation.FederationExecutionPolicy;
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.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 org.h2.jdbcx.JdbcDataSource;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

import java.sql.Connection;
import java.sql.Statement;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

/**
 * 四源和六源 JDBC 多阶段 Join 集成测试。
 */
public class JdbcSixSourceQueryEngineTest {

    private static final FederationExecutionPolicy SIX_SOURCE_POLICY =
        new FederationExecutionPolicy(6, 12, 2, 100_000,
            64L * 1024L * 1024L, 60_000);

    private final Map<SourceId, JdbcDataSource> dataSources = new LinkedHashMap<>();
    private final Map<String, FederationSourceBindingDefinition> bindings = new LinkedHashMap<>();
    private FederationSqlEngine engine;

    /**
     * 创建六个独立 H2 JDBC 数据源。
     */
    @Before
    public void setUp() throws Exception {
        for (int index = 1; index <= 6; index++) {
            SourceId sourceId = new SourceId("source-" + index);
            JdbcDataSource dataSource = dataSource("six-source-" + index);
            execute(dataSource,
                "DROP TABLE IF EXISTS SUBJECT_METRIC",
                "CREATE TABLE SUBJECT_METRIC (SUBJECT_ID INT PRIMARY KEY, METRIC_VALUE INT NOT NULL)",
                "INSERT INTO SUBJECT_METRIC VALUES (1, " + (index * 10)
                    + "), (2, " + (index * 100) + ")");
            dataSources.put(sourceId, dataSource);
            bindings.put("S" + index, FederationSourceBindingDefinition.of(sourceId, 1));
        }
        RuntimeFingerprint fingerprint = new RuntimeFingerprint(
            "H2", "2", "H2 JDBC Driver", "2", "1");
        engine = FederationSqlEngines.builder()
            .dataSourceResolver(definition -> FederationDataSourceHandles.shared(
                dataSources.get(definition.sourceId()), fingerprint))
            .adapter(new JdbcFederationSqlAdapterProvider())
            .federationExecutionPolicy(SIX_SOURCE_POLICY)
            .maximumPlanCacheEntries(32)
            .build();
        for (SourceId sourceId : dataSources.keySet()) {
            engine.sources().apply(definition(sourceId), SourceApplyOptions.prewarmNow());
        }
    }

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

    /**
     * 验证同一六源 Scope 中只引用四源时生成四个 Fragment。
     */
    @Test
    public void shouldCompileFourReferencedSources() {
        FederationSqlPlan plan = engine.compile(SqlCompileRequest.of(
            selectSql(4), scope(SIX_SOURCE_POLICY)));

        Assert.assertEquals(FederationQueryMode.FEDERATED, plan.queryMode());
        Assert.assertEquals(4, plan.referencedSources().size());
        Assert.assertEquals(4, plan.fragments().size());
        Assert.assertEquals(3, plan.joinOptimizations().size());
    }

    /**
     * 验证六个独立 JDBC 数据源可完成五阶段 Join 并返回稳定结果。
     */
    @Test
    public void shouldExecuteJoinAcrossSixSources() {
        String sql = selectSql(6);
        FederationSqlPlan plan = engine.compile(SqlCompileRequest.of(
            sql, scope(SIX_SOURCE_POLICY)));

        Assert.assertEquals(6, plan.referencedSources().size());
        Assert.assertEquals(6, plan.fragments().size());
        Assert.assertEquals(5, plan.joinOptimizations().size());

        try (FederationResultCursor cursor = engine.query(
            SqlQueryCommand.of(sql, scope(SIX_SOURCE_POLICY), List.of()))) {
            Assert.assertTrue(cursor.next());
            Assert.assertEquals(1, cursor.getObject(1));
            Assert.assertEquals(10, cursor.getObject(2));
            Assert.assertEquals(60, cursor.getObject(7));
            Assert.assertTrue(cursor.next());
            Assert.assertEquals(2, cursor.getObject(1));
            Assert.assertEquals(600, cursor.getObject(7));
            Assert.assertFalse(cursor.next());
            Assert.assertEquals(6, cursor.metrics().fragments().size());
        }
    }

    /**
     * 验证 Scope 收紧为四源时，六源 SQL 在编译期稳定拒绝。
     */
    @Test
    public void shouldRejectSixSourcesWhenScopeLimitIsFour() {
        FederationExecutionPolicy fourSourcePolicy = new FederationExecutionPolicy(
            4, 12, 2, 100_000, 64L * 1024L * 1024L, 60_000);

        try {
            engine.compile(SqlCompileRequest.of(selectSql(6), scope(fourSourcePolicy)));
            Assert.fail("expected source limit rejection");
        } catch (FederationSqlException exception) {
            Assert.assertEquals(FederationSqlErrorCode.FEDERATION_RESOURCE_LIMIT_EXCEEDED,
                exception.errorCode());
        }
    }

    private FederationQueryScopeDefinition scope(FederationExecutionPolicy policy) {
        return new FederationQueryScopeDefinition(
            "six-source-scope", 1, bindings, "S1", policy);
    }

    private static String selectSql(int sourceCount) {
        StringBuilder sql = new StringBuilder("SELECT s1.SUBJECT_ID");
        for (int index = 1; index <= sourceCount; index++) {
            sql.append(", s").append(index).append(".METRIC_VALUE");
        }
        sql.append(" FROM S1.APP.SUBJECT_METRIC s1");
        for (int index = 2; index <= sourceCount; index++) {
            sql.append(" JOIN S").append(index).append(".APP.SUBJECT_METRIC s")
                .append(index).append(" ON s").append(index)
                .append(".SUBJECT_ID = s1.SUBJECT_ID");
        }
        return sql.append(" ORDER BY s1.SUBJECT_ID").toString();
    }

    private static JdbcDataSource dataSource(String name) {
        JdbcDataSource dataSource = new JdbcDataSource();
        dataSource.setURL("jdbc:h2:mem:" + name + ";DB_CLOSE_DELAY=-1");
        dataSource.setUser("sa");
        dataSource.setPassword("");
        return dataSource;
    }

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