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

import com.easyagents.federation.sql.adapter.FederationStatisticsCollectionContext;
import com.easyagents.federation.sql.adapter.FederationStatisticsCollector;
import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import com.easyagents.federation.sql.federation.FederationColumnStatistics;
import com.easyagents.federation.sql.federation.FederationStatisticsSnapshot;
import com.easyagents.federation.sql.federation.FederationStatisticsStatus;
import com.easyagents.federation.sql.federation.FederationTableStatistics;
import com.easyagents.federation.sql.source.FederationSchemaDefinition;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Types;
import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * MySQL 与 PostgreSQL 的批量 JDBC 目录统计采集器。
 */
final class JdbcFederationStatisticsCollector implements FederationStatisticsCollector {

    private static final Logger LOG = LoggerFactory.getLogger(
        JdbcFederationStatisticsCollector.class
    );

    /**
     * 根据 JDBC 数据库产品分派内建统计采集逻辑。
     *
     * @param context 统计采集上下文
     * @return 表统计映射
     * @throws SQLException 目录或 JDBC 元数据读取失败
     */
    @Override
    public Map<FederationStatisticsSnapshot.TableKey, FederationTableStatistics> collect(
        FederationStatisticsCollectionContext context
    ) throws SQLException {
        String product = normalize(
            context.connection().getMetaData().getDatabaseProductName()
        );
        List<JdbcSchemaDefinition> schemas = jdbcSchemas(context);
        return switch (product) {
            case "mysql" -> collectMysql(context, schemas);
            case "postgresql" -> collectPostgresql(context, schemas);
            default -> Map.of();
        };
    }

    /**
     * 批量读取 MySQL INFORMATION_SCHEMA 表统计和主键。
     *
     * @param context 采集上下文
     * @param schemas JDBC Schema 映射
     * @return MySQL 表统计
     * @throws SQLException 目录读取失败
     */
    private Map<FederationStatisticsSnapshot.TableKey, FederationTableStatistics> collectMysql(
        FederationStatisticsCollectionContext context,
        List<JdbcSchemaDefinition> schemas
    ) throws SQLException {
        Map<FederationStatisticsSnapshot.TableKey, FederationTableStatistics> statistics =
            new LinkedHashMap<>();
        for (JdbcSchemaDefinition schema : schemas) {
            String catalog = textOr(schema.catalog(), context.connection().getCatalog());
            if (catalog == null || catalog.isBlank()) {
                continue;
            }
            Map<String, ColumnLayout> layouts = readColumnLayouts(
                context.connection().getMetaData(),
                catalog,
                schema.physicalSchema()
            );
            Map<String, List<String>> primaryKeys = readMysqlPrimaryKeys(
                context,
                catalog
            );
            String sql = "SELECT TABLE_NAME, TABLE_ROWS, AVG_ROW_LENGTH "
                + "FROM INFORMATION_SCHEMA.TABLES "
                + "WHERE TABLE_SCHEMA = ? AND TABLE_TYPE = 'BASE TABLE'";
            try (PreparedStatement statement = context.connection().prepareStatement(sql)) {
                statement.setQueryTimeout(context.queryTimeoutSeconds());
                statement.setString(1, catalog);
                try (ResultSet result = statement.executeQuery()) {
                    while (result.next()) {
                        String table = result.getString("TABLE_NAME");
                        ColumnLayout layout = layouts.getOrDefault(
                            normalize(table),
                            ColumnLayout.empty()
                        );
                        long averageWidth = result.getLong("AVG_ROW_LENGTH");
                        if (averageWidth <= 0L) {
                            averageWidth = layout.fallbackWidthBytes();
                        }
                        FederationStatisticsSnapshot.TableKey key =
                            new FederationStatisticsSnapshot.TableKey(
                                context.sourceDefinition().sourceId(),
                                schema.logicalName(),
                                table
                            );
                        statistics.put(key, new FederationTableStatistics(
                            Math.max(0D, result.getDouble("TABLE_ROWS")),
                            Math.max(1L, averageWidth),
                            context.collectedAt(),
                            "database-catalog:mysql",
                            Map.of(),
                            uniqueKey(primaryKeys.get(normalize(table))),
                            context.expiresAt(),
                            FederationStatisticsStatus.PARTIAL
                        ));
                    }
                }
            }
        }
        return Map.copyOf(statistics);
    }

    /**
     * 批量读取 PostgreSQL 表行数、列分布和主键统计。
     *
     * @param context 采集上下文
     * @param schemas JDBC Schema 映射
     * @return PostgreSQL 表统计
     * @throws SQLException 表级目录读取失败
     */
    private Map<FederationStatisticsSnapshot.TableKey, FederationTableStatistics>
        collectPostgresql(
            FederationStatisticsCollectionContext context,
            List<JdbcSchemaDefinition> schemas
        ) throws SQLException {
        Map<String, JdbcSchemaDefinition> schemasByPhysical = new LinkedHashMap<>();
        Map<String, ColumnLayout> layouts = new LinkedHashMap<>();
        for (JdbcSchemaDefinition schema : schemas) {
            String physical = textOr(schema.physicalSchema(), context.connection().getSchema());
            if (physical == null || physical.isBlank()) {
                physical = "public";
            }
            schemasByPhysical.putIfAbsent(normalize(physical), schema);
            Map<String, ColumnLayout> schemaLayouts = readColumnLayouts(
                context.connection().getMetaData(),
                schema.catalog(),
                physical
            );
            String resolvedPhysical = physical;
            schemaLayouts.forEach((table, layout) -> layouts.put(
                tableKey(resolvedPhysical, table),
                layout
            ));
        }
        if (schemasByPhysical.isEmpty()) {
            return Map.of();
        }
        Map<String, TableEstimate> estimates = readPostgresqlTableEstimates(
            context,
            schemasByPhysical.keySet()
        );
        Map<String, Map<String, FederationColumnStatistics>> columns;
        try {
            columns = readPostgresqlColumnStatistics(
                context,
                schemasByPhysical.keySet(),
                estimates
            );
        } catch (SQLException exception) {
            LOG.warn(
                "PostgreSQL column statistics are unavailable; retaining table estimates, sourceId={}",
                context.sourceDefinition().sourceId(),
                exception
            );
            columns = Map.of();
        }
        Map<String, List<String>> primaryKeys;
        try {
            primaryKeys = readPostgresqlPrimaryKeys(
                context,
                schemasByPhysical.keySet()
            );
        } catch (SQLException exception) {
            LOG.warn(
                "PostgreSQL primary-key statistics are unavailable, sourceId={}",
                context.sourceDefinition().sourceId(),
                exception
            );
            primaryKeys = Map.of();
        }

        Map<FederationStatisticsSnapshot.TableKey, FederationTableStatistics> statistics =
            new LinkedHashMap<>();
        for (Map.Entry<String, TableEstimate> entry : estimates.entrySet()) {
            TableEstimate estimate = entry.getValue();
            JdbcSchemaDefinition schema = schemasByPhysical.get(normalize(estimate.schema()));
            if (schema == null) {
                continue;
            }
            ColumnLayout layout = layouts.getOrDefault(entry.getKey(), ColumnLayout.empty());
            Map<String, FederationColumnStatistics> tableColumns = columns.getOrDefault(
                entry.getKey(),
                Map.of()
            );
            long averageWidth = Math.max(
                layout.fallbackWidthBytes(),
                averageColumnWidth(tableColumns)
            );
            boolean complete = !layout.columns().isEmpty()
                && containsAllIgnoreCase(tableColumns.keySet(), layout.columns());
            FederationStatisticsSnapshot.TableKey key =
                new FederationStatisticsSnapshot.TableKey(
                    context.sourceDefinition().sourceId(),
                    schema.logicalName(),
                    estimate.table()
                );
            statistics.put(key, new FederationTableStatistics(
                estimate.estimatedRows(),
                Math.max(1L, averageWidth),
                context.collectedAt(),
                "database-catalog:postgresql",
                tableColumns,
                uniqueKey(primaryKeys.get(entry.getKey())),
                context.expiresAt(),
                complete
                    ? FederationStatisticsStatus.COMPLETE
                    : FederationStatisticsStatus.PARTIAL
            ));
        }
        return Map.copyOf(statistics);
    }

    /**
     * 读取 Definition 中的 JDBC Schema 映射并拒绝不匹配的定义类型。
     *
     * @param context 采集上下文
     * @return JDBC Schema 定义
     */
    private List<JdbcSchemaDefinition> jdbcSchemas(
        FederationStatisticsCollectionContext context
    ) {
        List<JdbcSchemaDefinition> schemas = new ArrayList<>();
        for (FederationSchemaDefinition schema : context.sourceDefinition().schemas()) {
            if (!(schema instanceof JdbcSchemaDefinition jdbcSchema)) {
                throw new FederationSqlException(
                    FederationSqlErrorCode.INVALID_ARGUMENT,
                    "jdbc statistics collector requires JdbcSchemaDefinition"
                );
            }
            schemas.add(jdbcSchema);
        }
        return List.copyOf(schemas);
    }

    /**
     * 通过 JDBC 元数据按 Schema 批量读取字段布局。
     *
     * @param metadata JDBC 元数据
     * @param catalog 物理 Catalog
     * @param schema 物理 Schema
     * @return 按规范化表名索引的字段布局
     * @throws SQLException 元数据读取失败
     */
    private Map<String, ColumnLayout> readColumnLayouts(
        DatabaseMetaData metadata,
        String catalog,
        String schema
    ) throws SQLException {
        Map<String, MutableColumnLayout> layouts = new LinkedHashMap<>();
        try (ResultSet result = metadata.getColumns(catalog, schema, "%", "%")) {
            while (result.next()) {
                String table = normalize(result.getString("TABLE_NAME"));
                MutableColumnLayout layout = layouts.computeIfAbsent(
                    table,
                    ignored -> new MutableColumnLayout()
                );
                layout.columns.add(result.getString("COLUMN_NAME"));
                layout.fallbackWidthBytes = saturatedAdd(
                    layout.fallbackWidthBytes,
                    estimatedJdbcWidth(result.getInt("DATA_TYPE"))
                );
            }
        }
        Map<String, ColumnLayout> frozen = new LinkedHashMap<>();
        layouts.forEach((table, layout) -> frozen.put(
            table,
            new ColumnLayout(
                Math.max(1L, layout.fallbackWidthBytes),
                Set.copyOf(layout.columns)
            )
        ));
        return Map.copyOf(frozen);
    }

    /**
     * 一次查询一个 MySQL Catalog 的全部主键字段。
     *
     * @param context 采集上下文
     * @param catalog 物理 Catalog
     * @return 按规范化表名索引的有序主键
     * @throws SQLException 目录读取失败
     */
    private Map<String, List<String>> readMysqlPrimaryKeys(
        FederationStatisticsCollectionContext context,
        String catalog
    ) throws SQLException {
        String sql = "SELECT TABLE_NAME, COLUMN_NAME "
            + "FROM INFORMATION_SCHEMA.KEY_COLUMN_USAGE "
            + "WHERE TABLE_SCHEMA = ? AND CONSTRAINT_NAME = 'PRIMARY' "
            + "ORDER BY TABLE_NAME, ORDINAL_POSITION";
        Map<String, List<String>> keys = new LinkedHashMap<>();
        try (PreparedStatement statement = context.connection().prepareStatement(sql)) {
            statement.setQueryTimeout(context.queryTimeoutSeconds());
            statement.setString(1, catalog);
            try (ResultSet result = statement.executeQuery()) {
                while (result.next()) {
                    keys.computeIfAbsent(
                        normalize(result.getString("TABLE_NAME")),
                        ignored -> new ArrayList<>()
                    ).add(result.getString("COLUMN_NAME"));
                }
            }
        }
        return freezeLists(keys);
    }

    /**
     * 读取 PostgreSQL 表级近似行数。
     *
     * @param context 采集上下文
     * @param schemas 物理 Schema
     * @return 表级估算
     * @throws SQLException 目录读取失败
     */
    private Map<String, TableEstimate> readPostgresqlTableEstimates(
        FederationStatisticsCollectionContext context,
        Set<String> schemas
    ) throws SQLException {
        String sql = "SELECT n.nspname AS schema_name, c.relname AS table_name, "
            + "GREATEST(c.reltuples, 0)::double precision AS estimated_rows "
            + "FROM pg_catalog.pg_class c "
            + "JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace "
            + "WHERE c.relkind IN ('r', 'p') AND lower(n.nspname) IN ("
            + placeholders(schemas.size()) + ")";
        Map<String, TableEstimate> estimates = new LinkedHashMap<>();
        try (PreparedStatement statement = context.connection().prepareStatement(sql)) {
            statement.setQueryTimeout(context.queryTimeoutSeconds());
            bind(statement, schemas);
            try (ResultSet result = statement.executeQuery()) {
                while (result.next()) {
                    String schema = result.getString("schema_name");
                    String table = result.getString("table_name");
                    estimates.put(
                        tableKey(schema, table),
                        new TableEstimate(
                            schema,
                            table,
                            Math.max(0D, result.getDouble("estimated_rows"))
                        )
                    );
                }
            }
        }
        return Map.copyOf(estimates);
    }

    /**
     * 读取 PostgreSQL 列分布统计。
     *
     * @param context 采集上下文
     * @param schemas 物理 Schema
     * @param estimates 已读取的表级估算
     * @return 按物理表索引的列统计
     * @throws SQLException 目录读取失败
     */
    private Map<String, Map<String, FederationColumnStatistics>>
        readPostgresqlColumnStatistics(
            FederationStatisticsCollectionContext context,
            Set<String> schemas,
            Map<String, TableEstimate> estimates
        ) throws SQLException {
        String sql = "SELECT schemaname, tablename, attname, null_frac, n_distinct, avg_width "
            + "FROM pg_catalog.pg_stats WHERE lower(schemaname) IN ("
            + placeholders(schemas.size()) + ")";
        Map<String, Map<String, FederationColumnStatistics>> columns = new LinkedHashMap<>();
        try (PreparedStatement statement = context.connection().prepareStatement(sql)) {
            statement.setQueryTimeout(context.queryTimeoutSeconds());
            bind(statement, schemas);
            try (ResultSet result = statement.executeQuery()) {
                while (result.next()) {
                    String key = tableKey(
                        result.getString("schemaname"),
                        result.getString("tablename")
                    );
                    TableEstimate table = estimates.get(key);
                    if (table == null) {
                        continue;
                    }
                    double rawDistinct = result.getDouble("n_distinct");
                    double distinct = rawDistinct < 0D
                        ? Math.abs(rawDistinct) * table.estimatedRows()
                        : rawDistinct;
                    if (!Double.isFinite(distinct)) {
                        distinct = 0D;
                    }
                    columns.computeIfAbsent(key, ignored -> new LinkedHashMap<>()).put(
                        result.getString("attname"),
                        new FederationColumnStatistics(
                            Math.max(0D, distinct),
                            Math.max(0D, Math.min(1D, result.getDouble("null_frac"))),
                            Math.max(0L, result.getLong("avg_width"))
                        )
                    );
                }
            }
        }
        Map<String, Map<String, FederationColumnStatistics>> frozen = new LinkedHashMap<>();
        columns.forEach((table, values) -> frozen.put(table, Map.copyOf(values)));
        return Map.copyOf(frozen);
    }

    /**
     * 一次读取多个 PostgreSQL Schema 的主键字段。
     *
     * @param context 采集上下文
     * @param schemas 物理 Schema
     * @return 按物理表索引的主键字段
     * @throws SQLException 目录读取失败
     */
    private Map<String, List<String>> readPostgresqlPrimaryKeys(
        FederationStatisticsCollectionContext context,
        Set<String> schemas
    ) throws SQLException {
        String sql = "SELECT n.nspname AS schema_name, c.relname AS table_name, "
            + "a.attname AS column_name "
            + "FROM pg_catalog.pg_index i "
            + "JOIN pg_catalog.pg_class c ON c.oid = i.indrelid "
            + "JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace "
            + "JOIN pg_catalog.pg_attribute a "
            + "ON a.attrelid = c.oid AND a.attnum = ANY(i.indkey) "
            + "WHERE i.indisprimary AND lower(n.nspname) IN ("
            + placeholders(schemas.size()) + ") "
            + "ORDER BY n.nspname, c.relname, a.attnum";
        Map<String, List<String>> keys = new LinkedHashMap<>();
        try (PreparedStatement statement = context.connection().prepareStatement(sql)) {
            statement.setQueryTimeout(context.queryTimeoutSeconds());
            bind(statement, schemas);
            try (ResultSet result = statement.executeQuery()) {
                while (result.next()) {
                    keys.computeIfAbsent(
                        tableKey(
                            result.getString("schema_name"),
                            result.getString("table_name")
                        ),
                        ignored -> new ArrayList<>()
                    ).add(result.getString("column_name"));
                }
            }
        }
        return freezeLists(keys);
    }

    /**
     * 将可选主键转换为唯一键列表。
     *
     * @param primaryKey 主键字段
     * @return 零个或一个唯一键
     */
    private List<List<String>> uniqueKey(List<String> primaryKey) {
        return primaryKey == null || primaryKey.isEmpty()
            ? List.of()
            : List.of(List.copyOf(primaryKey));
    }

    /**
     * 汇总列平均宽度并防止 long 溢出。
     *
     * @param columns 列统计
     * @return 至少为 1 的平均宽度
     */
    private long averageColumnWidth(Map<String, FederationColumnStatistics> columns) {
        long width = 0L;
        for (FederationColumnStatistics column : columns.values()) {
            width = saturatedAdd(width, column.averageWidthBytes());
        }
        return Math.max(1L, width);
    }

    /**
     * 判断列统计是否覆盖全部字段。
     *
     * @param available 已有列统计名称
     * @param required JDBC 字段名称
     * @return 完整覆盖时为 true
     */
    private boolean containsAllIgnoreCase(Set<String> available, Set<String> required) {
        Set<String> normalized = new LinkedHashSet<>();
        available.forEach(value -> normalized.add(normalize(value)));
        return required.stream().map(this::normalize).allMatch(normalized::contains);
    }

    /**
     * 估算 JDBC 类型的保守内存宽度。
     *
     * @param jdbcType JDBC 类型
     * @return 估算字节数
     */
    private long estimatedJdbcWidth(int jdbcType) {
        return switch (jdbcType) {
            case Types.BOOLEAN, Types.BIT, Types.TINYINT -> 1L;
            case Types.SMALLINT -> 2L;
            case Types.INTEGER, Types.REAL, Types.FLOAT, Types.DATE -> 4L;
            case Types.BIGINT, Types.DOUBLE, Types.TIMESTAMP,
                Types.TIMESTAMP_WITH_TIMEZONE, Types.TIME,
                Types.TIME_WITH_TIMEZONE -> 8L;
            case Types.DECIMAL, Types.NUMERIC -> 16L;
            case Types.BINARY, Types.VARBINARY, Types.LONGVARBINARY,
                Types.BLOB, Types.CLOB, Types.NCLOB,
                Types.LONGVARCHAR, Types.LONGNVARCHAR -> 64L;
            default -> 32L;
        };
    }

    /**
     * 生成固定数量的 PreparedStatement 占位符。
     *
     * @param size 占位符数量
     * @return 逗号分隔占位符
     */
    private String placeholders(int size) {
        return String.join(", ", Collections.nCopies(size, "?"));
    }

    /**
     * 按稳定顺序绑定规范化 Schema。
     *
     * @param statement PreparedStatement
     * @param schemas Schema 集合
     * @throws SQLException 参数绑定失败
     */
    private void bind(PreparedStatement statement, Set<String> schemas) throws SQLException {
        int index = 1;
        for (String schema : schemas) {
            statement.setString(index++, normalize(schema));
        }
    }

    /**
     * 冻结可变列表映射。
     *
     * @param source 可变列表映射
     * @return 不可变列表映射
     */
    private Map<String, List<String>> freezeLists(Map<String, List<String>> source) {
        Map<String, List<String>> frozen = new LinkedHashMap<>();
        source.forEach((key, value) -> frozen.put(key, List.copyOf(value)));
        return Map.copyOf(frozen);
    }

    /**
     * 生成大小写不敏感的物理表索引键。
     *
     * @param schema 物理 Schema
     * @param table 物理表
     * @return 稳定索引键
     */
    private String tableKey(String schema, String table) {
        return normalize(schema) + '\u0000' + normalize(table);
    }

    /**
     * 返回首个非空文本。
     *
     * @param primary 首选值
     * @param fallback 备用值
     * @return 可空结果
     */
    private String textOr(String primary, String fallback) {
        return primary == null || primary.isBlank() ? fallback : primary;
    }

    /**
     * 规范化数据库产品名或标识符。
     *
     * @param value 原始值
     * @return 小写非空值
     */
    private String normalize(String value) {
        return value == null ? "" : value.trim().toLowerCase(Locale.ROOT);
    }

    /**
     * 饱和 long 加法。
     *
     * @param left 左值
     * @param right 右值
     * @return 不溢出的和
     */
    private long saturatedAdd(long left, long right) {
        return Long.MAX_VALUE - left < right ? Long.MAX_VALUE : left + right;
    }

    /**
     * 单表字段布局。
     *
     * @param fallbackWidthBytes JDBC 类型估算行宽
     * @param columns 字段名称
     */
    private record ColumnLayout(long fallbackWidthBytes, Set<String> columns) {

        /**
         * 创建空布局。
         *
         * @return 保守空布局
         */
        private static ColumnLayout empty() {
            return new ColumnLayout(1L, Set.of());
        }
    }

    /** 可变字段布局构造器。 */
    private static final class MutableColumnLayout {
        private long fallbackWidthBytes;
        private final Set<String> columns = new LinkedHashSet<>();
    }

    /**
     * PostgreSQL 表级估算。
     *
     * @param schema 物理 Schema
     * @param table 物理表
     * @param estimatedRows 估算行数
     */
    private record TableEstimate(String schema, String table, double estimatedRows) {
    }
}
