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

import com.easyagents.federation.sql.adapter.FederationStatisticsCollectionContext;
import com.easyagents.federation.sql.api.FederationSqlEngine;
import com.easyagents.federation.sql.api.FederationSqlEngines;
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.federation.FederationExecutionPolicy;
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.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 java.sql.Connection;
import java.sql.DriverManager;
import java.time.Duration;
import java.time.Instant;
import java.util.List;
import java.util.Map;
import javax.sql.DataSource;
import org.postgresql.ds.PGSimpleDataSource;
import org.junit.Assert;
import org.junit.Assume;
import org.junit.Test;

/**
 * 本机 MySQL 与 PostgreSQL 的 Adapter 内建统计采集验证。
 */
public class JdbcFederationStatisticsIntegrationTest {

    private static final Duration STATISTICS_TTL = Duration.ofMinutes(30);
    private static final SourceId MYSQL_SOURCE = new SourceId("mysql-statistics");
    private static final SourceId POSTGRESQL_SOURCE = new SourceId(
        "postgresql-statistics"
    );

    /**
     * 验证 MySQL 目录统计可由 JDBC Adapter 自动采集。
     *
     * @throws Exception JDBC 连接或目录读取失败
     */
    @Test
    public void shouldCollectMysqlStatistics() throws Exception {
        Assume.assumeTrue(Boolean.getBoolean("federation.integration.enabled"));
        String database = System.getProperty(
            "federation.mysql.database",
            "data-sheet"
        );
        String url = "jdbc:mysql://"
            + System.getProperty("federation.mysql.host", "127.0.0.1")
            + ':' + Integer.getInteger("federation.mysql.port", 33306)
            + '/' + database
            + "?useSSL=false&allowPublicKeyRetrieval=true&serverTimezone=Asia/Shanghai";
        FederationSourceDefinition definition = definition(
            "mysql-statistics",
            "MAIN",
            database,
            null
        );

        try (Connection connection = DriverManager.getConnection(
            url,
            System.getProperty("federation.mysql.username", "root"),
            System.getProperty("federation.mysql.password", "root")
        )) {
            Map<FederationStatisticsSnapshot.TableKey, FederationTableStatistics> statistics =
                collect(definition, connection);

            FederationTableStatistics outlet = statistics.get(
                new FederationStatisticsSnapshot.TableKey(
                    definition.sourceId(),
                    "MAIN",
                    "outlet"
                )
            );
            Assert.assertNotNull(outlet);
            assertUsable(outlet, "database-catalog:mysql");
        }
    }

    /**
     * 验证 PostgreSQL 目录统计可由 JDBC Adapter 自动采集。
     *
     * @throws Exception JDBC 连接或目录读取失败
     */
    @Test
    public void shouldCollectPostgresqlStatistics() throws Exception {
        Assume.assumeTrue(Boolean.getBoolean("federation.integration.enabled"));
        String database = System.getProperty(
            "federation.pg.database",
            "harmony_adapter"
        );
        String url = "jdbc:postgresql://"
            + System.getProperty("federation.pg.host", "127.0.0.1")
            + ':' + Integer.getInteger("federation.pg.port", 54329)
            + '/' + database;
        FederationSourceDefinition definition = definition(
            "postgresql-statistics",
            "MAIN",
            database,
            "public"
        );

        try (Connection connection = DriverManager.getConnection(
            url,
            System.getProperty("federation.pg.username", "harmony"),
            System.getProperty("federation.pg.password", "harmony")
        )) {
            Map<FederationStatisticsSnapshot.TableKey, FederationTableStatistics> statistics =
                collect(definition, connection);

            Assert.assertFalse(statistics.isEmpty());
            Assert.assertTrue(statistics.keySet().stream().allMatch(key ->
                definition.sourceId().equals(key.sourceId())
                    && "main".equals(key.schema())
            ));
            statistics.values().forEach(value ->
                assertUsable(value, "database-catalog:postgresql")
            );
        }
    }

    /**
     * 验证 Engine 默认启用 Adapter 统计，并将结果交给 Explain 成本估算。
     */
    @Test
    public void shouldExposeAutomaticallyCollectedStatisticsThroughExplain() {
        Assume.assumeTrue(Boolean.getBoolean("federation.integration.enabled"));
        String mysqlDatabase = System.getProperty(
            "federation.mysql.database",
            "data-sheet"
        );
        String postgresqlDatabase = System.getProperty(
            "federation.pg.database",
            "harmony_adapter"
        );
        FederationSourceDefinition mysqlDefinition = definition(
            MYSQL_SOURCE.value(),
            "main",
            mysqlDatabase,
            null
        );
        FederationSourceDefinition postgresqlDefinition = definition(
            POSTGRESQL_SOURCE.value(),
            "main",
            postgresqlDatabase,
            "public"
        );
        Map<SourceId, DataSource> dataSources = Map.of(
            MYSQL_SOURCE,
            mysqlDataSource(mysqlDatabase),
            POSTGRESQL_SOURCE,
            postgresqlDataSource(postgresqlDatabase)
        );

        try (FederationSqlEngine engine = FederationSqlEngines.builder()
            .dataSourceResolver(definition -> FederationDataSourceHandles.shared(
                dataSources.get(definition.sourceId()),
                new RuntimeFingerprint("test", "1", "jdbc", "1", "1")
            ))
            .build()) {
            engine.sources().apply(mysqlDefinition, SourceApplyOptions.prewarmNow());
            engine.sources().apply(
                postgresqlDefinition,
                SourceApplyOptions.prewarmNow()
            );
            FederationQueryScopeDefinition scope = FederationQueryScopeDefinition.virtual(
                "automatic-statistics",
                1,
                Map.of(
                    "mysql",
                    FederationSourceBindingDefinition.of(MYSQL_SOURCE, 1),
                    "pg",
                    FederationSourceBindingDefinition.of(POSTGRESQL_SOURCE, 1)
                ),
                "mysql",
                FederationExecutionPolicy.basic()
            );

            SqlExplainResult mysqlExplain = explain(
                engine,
                scope,
                "SELECT * FROM mysql.main.outlet"
            );
            SqlExplainResult postgresqlExplain = explain(
                engine,
                scope,
                "SELECT * FROM pg.main.artifact"
            );

            assertExplainStatistics(mysqlExplain, "database-catalog:mysql");
            assertExplainStatistics(
                postgresqlExplain,
                "database-catalog:postgresql"
            );
        }
    }

    /**
     * 使用 Adapter 统计采集器读取当前连接。
     *
     * @param definition 物理源定义
     * @param connection JDBC 连接
     * @return 不可变表统计
     * @throws Exception 目录读取失败
     */
    private Map<FederationStatisticsSnapshot.TableKey, FederationTableStatistics> collect(
        FederationSourceDefinition definition,
        Connection connection
    ) throws Exception {
        Instant collectedAt = Instant.now();
        return new JdbcFederationStatisticsCollector().collect(
            new FederationStatisticsCollectionContext(
                definition,
                connection,
                collectedAt,
                collectedAt.plus(STATISTICS_TTL),
                5
            )
        );
    }

    /**
     * 创建单 Schema JDBC 数据源定义。
     *
     * @param sourceId 物理源标识
     * @param logicalSchema 逻辑 Schema
     * @param catalog 物理 Catalog
     * @param physicalSchema 物理 Schema
     * @return 数据源定义
     */
    private FederationSourceDefinition definition(
        String sourceId,
        String logicalSchema,
        String catalog,
        String physicalSchema
    ) {
        return new FederationSourceDefinition(
            new SourceId(sourceId),
            1,
            JdbcFederationSqlAdapterProvider.ADAPTER_ID,
            List.of(new JdbcSchemaDefinition(
                logicalSchema,
                catalog,
                physicalSchema
            )),
            Map.of()
        );
    }

    /**
     * 创建本机 MySQL 测试 DataSource。
     *
     * @param database 数据库名称
     * @return MySQL DataSource
     */
    private DataSource mysqlDataSource(String database) {
        MysqlDataSource dataSource = new MysqlDataSource();
        dataSource.setUrl(
            "jdbc:mysql://"
                + System.getProperty("federation.mysql.host", "127.0.0.1")
                + ':' + Integer.getInteger("federation.mysql.port", 33306)
                + '/' + database
                + "?useSSL=false&allowPublicKeyRetrieval=true&serverTimezone=Asia/Shanghai"
        );
        dataSource.setUser(System.getProperty("federation.mysql.username", "root"));
        dataSource.setPassword(System.getProperty("federation.mysql.password", "root"));
        return dataSource;
    }

    /**
     * 创建本机 PostgreSQL 测试 DataSource。
     *
     * @param database 数据库名称
     * @return PostgreSQL DataSource
     */
    private DataSource postgresqlDataSource(String database) {
        PGSimpleDataSource dataSource = new PGSimpleDataSource();
        dataSource.setServerNames(new String[]{
            System.getProperty("federation.pg.host", "127.0.0.1")
        });
        dataSource.setPortNumbers(new int[]{
            Integer.getInteger("federation.pg.port", 54329)
        });
        dataSource.setDatabaseName(database);
        dataSource.setUser(System.getProperty("federation.pg.username", "harmony"));
        dataSource.setPassword(System.getProperty("federation.pg.password", "harmony"));
        return dataSource;
    }

    /**
     * 执行逻辑 Explain。
     *
     * @param engine 联邦 SQL Engine
     * @param scope 查询范围
     * @param sql SQL
     * @return Explain 结果
     */
    private SqlExplainResult explain(
        FederationSqlEngine engine,
        FederationQueryScopeDefinition scope,
        String sql
    ) {
        return engine.explain(new SqlExplainRequest(
            SqlCompileRequest.of(sql, scope),
            SqlExplainLevel.LOGICAL
        ));
    }

    /**
     * 断言 Explain 已使用自动采集的数据库统计。
     *
     * @param explain Explain 结果
     * @param source 预期统计来源
     */
    private void assertExplainStatistics(SqlExplainResult explain, String source) {
        Assert.assertEquals(1, explain.fragments().size());
        Assert.assertFalse(explain.fragments().get(0).costEstimate().statisticsMissing());
        Assert.assertTrue(
            explain.fragments().get(0).costEstimate().statisticsSource().contains(source)
        );
        Assert.assertTrue(
            explain.fragments().get(0).costEstimate().estimatedRowWidthBytes() > 0L
        );
    }

    /**
     * 断言采集结果包含优化器可使用的基础统计。
     *
     * @param statistics 表统计
     * @param source 预期统计来源
     */
    private void assertUsable(FederationTableStatistics statistics, String source) {
        Assert.assertTrue(statistics.estimatedRows() >= 0D);
        Assert.assertTrue(statistics.averageRowWidthBytes() > 0L);
        Assert.assertEquals(source, statistics.source());
        Assert.assertNotNull(statistics.collectedAt());
        Assert.assertTrue(statistics.expiresAt().isAfter(statistics.collectedAt()));
    }
}
