package com.easyagents.federation.sql.runtime;

import com.easyagents.federation.sql.federation.FederationColumnStatistics;
import com.easyagents.federation.sql.federation.FederationStatisticsStatus;
import com.easyagents.federation.sql.federation.FederationTableStatistics;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import org.apache.calcite.plan.RelOptCost;
import org.apache.calcite.plan.RelOptCostImpl;
import org.apache.calcite.plan.RelOptCluster;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.core.Join;
import org.apache.calcite.rel.metadata.BuiltInMetadata;
import org.apache.calcite.rel.metadata.ChainedRelMetadataProvider;
import org.apache.calcite.rel.metadata.MetadataDef;
import org.apache.calcite.rel.metadata.MetadataHandler;
import org.apache.calcite.rel.metadata.RelMetadataProvider;
import org.apache.calcite.rel.metadata.RelMetadataQuery;
import org.apache.calcite.rel.metadata.ReflectiveRelMetadataProvider;
import org.apache.calcite.rex.RexCall;
import org.apache.calcite.rex.RexInputRef;
import org.apache.calcite.rex.RexLiteral;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.rex.RexUtil;
import org.apache.calcite.sql.SqlKind;
import org.apache.calcite.util.ImmutableBitSet;

/**
 * 将调用方冻结的表列统计接入 Calcite 元数据查询链路。
 *
 * <p>该类型仅因 Calcite 元数据代码生成要求公开可访问，调用方无需直接使用。</p>
 */
public final class FederationStatisticsMetadata {

    private static final RelMetadataProvider SOURCE = ChainedRelMetadataProvider.of(List.of(
        ReflectiveRelMetadataProvider.reflectiveSource(
            new CumulativeCostHandler(),
            BuiltInMetadata.CumulativeCost.Handler.class
        ),
        ReflectiveRelMetadataProvider.reflectiveSource(
            new DistinctRowCountHandler(),
            BuiltInMetadata.DistinctRowCount.Handler.class
        ),
        ReflectiveRelMetadataProvider.reflectiveSource(
            new SelectivityHandler(),
            BuiltInMetadata.Selectivity.Handler.class
        ),
        ReflectiveRelMetadataProvider.reflectiveSource(
            new UniqueKeysHandler(),
            BuiltInMetadata.UniqueKeys.Handler.class
        ),
        ReflectiveRelMetadataProvider.reflectiveSource(
            new SizeHandler(),
            BuiltInMetadata.Size.Handler.class
        )
    ));

    /**
     * 禁止实例化静态元数据桥接器。
     */
    private FederationStatisticsMetadata() { }

    /**
     * 将联邦统计提供器置于现有元数据链路之前，并清理已经计算的元数据。
     *
     * @param cluster 关系表达式集群
     */
    static void install(RelOptCluster cluster) {
        cluster.setMetadataProvider(ChainedRelMetadataProvider.of(List.of(
            SOURCE,
            cluster.getMetadataProvider()
        )));
        cluster.invalidateMetadataQuery();
    }

    /**
     * 按字段序号查找列统计，并兼容数据库返回的字段名大小写差异。
     *
     * @param scan 表扫描
     * @param index 字段序号
     * @return 列统计；字段或统计不存在时返回 null
     */
    private static FederationColumnStatistics columnStatistics(
        FederationStatisticsTableScan scan,
        int index
    ) {
        if (index < 0 || index >= scan.getRowType().getFieldCount()) {
            return null;
        }
        String fieldName = scan.getRowType().getFieldList().get(index).getName();
        Map<String, FederationColumnStatistics> columns = scan.statistics().columns();
        FederationColumnStatistics direct = columns.get(fieldName);
        if (direct != null) {
            return direct;
        }
        String normalized = fieldName.toLowerCase(Locale.ROOT);
        return columns.entrySet().stream()
            .filter(entry -> entry.getKey().toLowerCase(Locale.ROOT).equals(normalized))
            .map(Map.Entry::getValue)
            .findFirst()
            .orElse(null);
    }

    /**
     * 提取输入字段序号，同时剥离不改变字段来源的 CAST。
     *
     * @param node 表达式节点
     * @return 输入字段序号；非字段引用时返回 -1
     */
    private static int inputIndex(RexNode node) {
        RexNode current = node;
        while (current instanceof RexCall call
            && call.getKind() == SqlKind.CAST
            && RexUtil.isLosslessCast(call)) {
            current = call.getOperands().get(0);
        }
        return current instanceof RexInputRef inputRef ? inputRef.getIndex() : -1;
    }

    /**
     * 将概率值限制在有效区间。
     *
     * @param value 原始概率
     * @return 位于 0 到 1 之间的概率
     */
    private static double clampProbability(double value) {
        return Math.max(0D, Math.min(1D, value));
    }

    /**
     * 将成本值限制在可比较的有限区间。
     *
     * @param value 原始成本
     * @return 有限非负成本
     */
    private static double finiteCost(double value) {
        if (Double.isNaN(value) || value <= 0D) {
            return 0D;
        }
        return Math.min(value, Double.MAX_VALUE / 16D);
    }

    /**
     * 安全累加成本，避免极端统计导致无穷值破坏计划比较。
     *
     * @param left 左成本
     * @param right 右成本
     * @return 饱和后的成本
     */
    private static double addCost(double left, double right) {
        double normalizedLeft = finiteCost(left);
        double normalizedRight = finiteCost(right);
        double maximum = Double.MAX_VALUE / 16D;
        return normalizedLeft >= maximum - normalizedRight
            ? maximum
            : normalizedLeft + normalizedRight;
    }

    /**
     * 安全计算乘积成本。
     *
     * @param left 左因子
     * @param right 右因子
     * @return 饱和后的成本
     */
    private static double multiplyCost(double left, double right) {
        double normalizedLeft = finiteCost(left);
        double normalizedRight = finiteCost(right);
        double maximum = Double.MAX_VALUE / 16D;
        if (normalizedLeft == 0D || normalizedRight == 0D) {
            return 0D;
        }
        return normalizedLeft >= maximum / normalizedRight
            ? maximum
            : normalizedLeft * normalizedRight;
    }

    /**
     * 将 Calcite 多维成本折算成稳定标量。
     *
     * @param cost 原始成本
     * @return 非负有限成本
     */
    private static double scalarCost(RelOptCost cost) {
        if (cost == null || cost.isInfinite()) {
            return Double.MAX_VALUE / 16D;
        }
        return addCost(addCost(cost.getRows(), cost.getCpu()), cost.getIo());
    }

    /**
     * 为联邦 Join 搜索提供包含数据搬运和本地构建开销的累计成本。
     */
    public static final class CumulativeCostHandler
        implements MetadataHandler<BuiltInMetadata.CumulativeCost> {

        private static final double HASH_BUILD_WEIGHT = 2D;

        /**
         * 创建累计联邦成本处理器。
         */
        private CumulativeCostHandler() { }

        /** {@inheritDoc} */
        @Override
        public MetadataDef<BuiltInMetadata.CumulativeCost> getDef() {
            return BuiltInMetadata.CumulativeCost.DEF;
        }

        /**
         * 计算通用节点的自身成本和全部输入累计成本。
         *
         * @param node 关系节点
         * @param metadataQuery 元数据查询
         * @return 累计成本
         */
        public RelOptCost getCumulativeCost(
            RelNode node,
            RelMetadataQuery metadataQuery
        ) {
            double total = scalarCost(metadataQuery.getNonCumulativeCost(node));
            for (RelNode input : node.getInputs()) {
                total = addCost(total, scalarCost(metadataQuery.getCumulativeCost(input)));
            }
            return new RelOptCostImpl(total);
        }

        /**
         * 使用冻结行数与行宽计算物理源数据搬运成本。
         *
         * @param scan 表扫描
         * @param metadataQuery 元数据查询
         * @return 扫描累计成本
         */
        public RelOptCost getCumulativeCost(
            FederationStatisticsTableScan scan,
            RelMetadataQuery metadataQuery
        ) {
            FederationTableStatistics statistics = scan.statistics();
            double transfer = multiplyCost(
                statistics.estimatedRows(),
                statistics.averageRowWidthBytes()
            );
            return new RelOptCostImpl(addCost(statistics.estimatedRows(), transfer));
        }

        /**
         * 计算 Join 输入、输出与右侧 Hash 构建的累计成本。
         *
         * @param join Join 节点
         * @param metadataQuery 元数据查询
         * @return Join 累计成本
         */
        public RelOptCost getCumulativeCost(
            Join join,
            RelMetadataQuery metadataQuery
        ) {
            double leftCost = scalarCost(metadataQuery.getCumulativeCost(join.getLeft()));
            double rightCost = scalarCost(metadataQuery.getCumulativeCost(join.getRight()));
            double outputRows = finiteCost(metadataQuery.getRowCount(join));
            Double averageRowSize = metadataQuery.getAverageRowSize(join);
            double outputWidth = averageRowSize == null ? 16D : finiteCost(averageRowSize);
            double outputCost = addCost(outputRows, multiplyCost(outputRows, outputWidth));
            double rightRows = finiteCost(metadataQuery.getRowCount(join.getRight()));
            Double rightRowSize = metadataQuery.getAverageRowSize(join.getRight());
            double rightWidth = rightRowSize == null ? 16D : finiteCost(rightRowSize);
            double buildCost = multiplyCost(
                multiplyCost(rightRows, rightWidth),
                HASH_BUILD_WEIGHT
            );
            return new RelOptCostImpl(addCost(
                addCost(leftCost, rightCost),
                addCost(outputCost, buildCost)
            ));
        }
    }

    /**
     * 使用冻结列统计估算基础谓词选择率。
     *
     * @param scan 表扫描
     * @param predicate 谓词表达式
     * @return 选择率；无法可靠估算时返回 null
     */
    private static Double predicateSelectivity(
        FederationStatisticsTableScan scan,
        RexNode predicate
    ) {
        if (predicate == null) {
            return 1D;
        }
        if (!(predicate instanceof RexCall call)) {
            return null;
        }
        return switch (call.getKind()) {
            case AND -> combineAnd(scan, call.getOperands());
            case OR -> combineOr(scan, call.getOperands());
            case EQUALS -> equalitySelectivity(scan, call, false);
            case IS_NOT_DISTINCT_FROM -> equalitySelectivity(scan, call, true);
            case IS_NULL -> nullSelectivity(scan, call, true);
            case IS_NOT_NULL -> nullSelectivity(scan, call, false);
            default -> null;
        };
    }

    /**
     * 合并 AND 子谓词选择率。
     *
     * @param scan 表扫描
     * @param operands 子谓词
     * @return 合并选择率；任一子谓词不可估算时返回 null
     */
    private static Double combineAnd(
        FederationStatisticsTableScan scan,
        List<RexNode> operands
    ) {
        double result = 1D;
        for (RexNode operand : operands) {
            Double selectivity = predicateSelectivity(scan, operand);
            if (selectivity == null) {
                return null;
            }
            result *= selectivity;
        }
        return clampProbability(result);
    }

    /**
     * 合并 OR 子谓词选择率。
     *
     * @param scan 表扫描
     * @param operands 子谓词
     * @return 合并选择率；任一子谓词不可估算时返回 null
     */
    private static Double combineOr(
        FederationStatisticsTableScan scan,
        List<RexNode> operands
    ) {
        double noneSelected = 1D;
        for (RexNode operand : operands) {
            Double selectivity = predicateSelectivity(scan, operand);
            if (selectivity == null) {
                return null;
            }
            noneSelected *= 1D - selectivity;
        }
        return clampProbability(1D - noneSelected);
    }

    /**
     * 根据参与等值比较列的基数和空值比例估算选择率。
     *
     * @param scan 表扫描
     * @param call 等值表达式
     * @param nullSafe 是否采用 NULL-safe 等值语义
     * @return 选择率；统计不足时返回 null
     */
    private static Double equalitySelectivity(
        FederationStatisticsTableScan scan,
        RexCall call,
        boolean nullSafe
    ) {
        if (call.getOperands().size() != 2) {
            return null;
        }
        int leftIndex = inputIndex(call.getOperands().get(0));
        int rightIndex = inputIndex(call.getOperands().get(1));
        if (leftIndex < 0 && rightIndex < 0) {
            return null;
        }
        FederationColumnStatistics left = columnStatistics(scan, leftIndex);
        FederationColumnStatistics right = columnStatistics(scan, rightIndex);
        if (nullSafe && isNullLiteral(call.getOperands().get(0)) && right != null) {
            return right.nullFraction();
        }
        if (nullSafe && isNullLiteral(call.getOperands().get(1)) && left != null) {
            return left.nullFraction();
        }
        if (leftIndex >= 0 && leftIndex == rightIndex && left != null) {
            // 同一字段与自身比较时，普通等值仅排除 NULL，NULL-safe 等值恒为真。
            return nullSafe ? 1D : 1D - left.nullFraction();
        }
        double distinctCount = Math.max(
            left == null ? 0D : left.distinctCount(),
            right == null ? 0D : right.distinctCount()
        );
        if (distinctCount <= 0D) {
            return null;
        }
        double nonNullFraction = 1D;
        if (left != null) {
            nonNullFraction *= 1D - left.nullFraction();
        }
        if (right != null) {
            nonNullFraction *= 1D - right.nullFraction();
        }
        double selectivity = nonNullFraction / distinctCount;
        if (nullSafe && left != null && right != null) {
            // NULL-safe 等值还会命中两侧同时为 NULL 的行。
            selectivity += left.nullFraction() * right.nullFraction();
        }
        return clampProbability(selectivity);
    }

    /**
     * 判断表达式是否为 NULL 字面量，仅剥离不改变 NULL 语义的 CAST。
     *
     * @param node 表达式节点
     * @return 是否为 NULL 字面量
     */
    private static boolean isNullLiteral(RexNode node) {
        RexNode current = node;
        while (current instanceof RexCall call && call.getKind() == SqlKind.CAST) {
            current = call.getOperands().get(0);
        }
        return current instanceof RexLiteral literal && literal.isNull();
    }

    /**
     * 估算 IS NULL 或 IS NOT NULL 谓词选择率。
     *
     * @param scan 表扫描
     * @param call 空值判断表达式
     * @param nullOnly true 表示 IS NULL，false 表示 IS NOT NULL
     * @return 选择率；统计不足时返回 null
     */
    private static Double nullSelectivity(
        FederationStatisticsTableScan scan,
        RexCall call,
        boolean nullOnly
    ) {
        if (call.getOperands().size() != 1) {
            return null;
        }
        FederationColumnStatistics column = columnStatistics(
            scan,
            inputIndex(call.getOperands().get(0))
        );
        if (column == null) {
            return null;
        }
        return nullOnly ? column.nullFraction() : 1D - column.nullFraction();
    }

    /**
     * 基于列基数估算分组后的不同值数量。
     */
    public static final class DistinctRowCountHandler
        implements MetadataHandler<BuiltInMetadata.DistinctRowCount> {

        /**
         * 创建不同值行数元数据处理器。
         */
        private DistinctRowCountHandler() { }

        /** {@inheritDoc} */
        @Override
        public MetadataDef<BuiltInMetadata.DistinctRowCount> getDef() {
            return BuiltInMetadata.DistinctRowCount.DEF;
        }

        /**
         * 返回冻结统计中的不同值估算。
         *
         * @param scan 表扫描
         * @param metadataQuery 元数据查询
         * @param groupKey 分组列
         * @param predicate 可选过滤条件
         * @return 不同值估算；统计不足时返回 null 交给默认提供器
         */
        public Double getDistinctRowCount(
            FederationStatisticsTableScan scan,
            RelMetadataQuery metadataQuery,
            ImmutableBitSet groupKey,
            RexNode predicate
        ) {
            if (groupKey.isEmpty()) {
                return null;
            }
            double distinctRows = 1D;
            for (int index : groupKey) {
                FederationColumnStatistics column = columnStatistics(scan, index);
                if (column == null || column.distinctCount() <= 0D) {
                    return null;
                }
                distinctRows = Math.min(
                    scan.statistics().estimatedRows(),
                    distinctRows * column.distinctCount()
                );
            }
            Double selectivity = predicateSelectivity(scan, predicate);
            double selectedRows = selectivity == null
                ? scan.statistics().estimatedRows()
                : scan.statistics().estimatedRows() * selectivity;
            return Math.min(distinctRows, selectedRows);
        }
    }

    /**
     * 基于列基数和空值比例估算基础谓词选择率。
     */
    public static final class SelectivityHandler
        implements MetadataHandler<BuiltInMetadata.Selectivity> {

        /**
         * 创建选择率元数据处理器。
         */
        private SelectivityHandler() { }

        /** {@inheritDoc} */
        @Override
        public MetadataDef<BuiltInMetadata.Selectivity> getDef() {
            return BuiltInMetadata.Selectivity.DEF;
        }

        /**
         * 返回基础谓词的选择率。
         *
         * @param scan 表扫描
         * @param metadataQuery 元数据查询
         * @param predicate 过滤条件
         * @return 选择率；无法可靠估算时返回 null 交给默认提供器
         */
        public Double getSelectivity(
            FederationStatisticsTableScan scan,
            RelMetadataQuery metadataQuery,
            RexNode predicate
        ) {
            return predicateSelectivity(scan, predicate);
        }
    }

    /**
     * 将调用方声明的唯一键转换为 Calcite 列位集合。
     */
    public static final class UniqueKeysHandler
        implements MetadataHandler<BuiltInMetadata.UniqueKeys> {

        /**
         * 创建唯一键元数据处理器。
         */
        private UniqueKeysHandler() { }

        /** {@inheritDoc} */
        @Override
        public MetadataDef<BuiltInMetadata.UniqueKeys> getDef() {
            return BuiltInMetadata.UniqueKeys.DEF;
        }

        /**
         * 返回冻结统计中的唯一键。
         *
         * @param scan 表扫描
         * @param metadataQuery 元数据查询
         * @param ignoreNulls 是否忽略空值
         * @return 唯一键集合；部分统计且无键时返回 null
         */
        public Set<ImmutableBitSet> getUniqueKeys(
            FederationStatisticsTableScan scan,
            RelMetadataQuery metadataQuery,
            boolean ignoreNulls
        ) {
            FederationTableStatistics statistics = scan.statistics();
            Set<ImmutableBitSet> keys = new HashSet<>();
            List<String> fields = scan.getRowType().getFieldNames();
            for (List<String> key : statistics.uniqueKeys()) {
                ImmutableBitSet.Builder builder = ImmutableBitSet.builder();
                boolean complete = true;
                for (String columnName : key) {
                    int index = findField(fields, columnName);
                    if (index < 0) {
                        complete = false;
                        break;
                    }
                    builder.set(index);
                }
                if (complete && !key.isEmpty()) {
                    keys.add(builder.build());
                }
            }
            if (keys.isEmpty() && statistics.status() != FederationStatisticsStatus.COMPLETE) {
                return null;
            }
            return Set.copyOf(keys);
        }

        /**
         * 按不区分大小写的方式定位字段。
         *
         * @param fields 字段名称列表
         * @param expected 待查找字段
         * @return 字段序号；未找到时返回 -1
         */
        private static int findField(List<String> fields, String expected) {
            for (int index = 0; index < fields.size(); index++) {
                if (fields.get(index).equalsIgnoreCase(expected)) {
                    return index;
                }
            }
            return -1;
        }
    }

    /**
     * 将平均行宽与列宽接入 Calcite Size 元数据。
     */
    public static final class SizeHandler implements MetadataHandler<BuiltInMetadata.Size> {

        /**
         * 创建行列宽度元数据处理器。
         */
        private SizeHandler() { }

        /** {@inheritDoc} */
        @Override
        public MetadataDef<BuiltInMetadata.Size> getDef() {
            return BuiltInMetadata.Size.DEF;
        }

        /**
         * 返回平均行宽。
         *
         * @param scan 表扫描
         * @param metadataQuery 元数据查询
         * @return 平均行宽字节数
         */
        public Double averageRowSize(
            FederationStatisticsTableScan scan,
            RelMetadataQuery metadataQuery
        ) {
            return (double) scan.statistics().averageRowWidthBytes();
        }

        /**
         * 返回各列平均宽度。
         *
         * @param scan 表扫描
         * @param metadataQuery 元数据查询
         * @return 列宽列表；没有任何列宽统计时返回 null
         */
        public List<Double> averageColumnSizes(
            FederationStatisticsTableScan scan,
            RelMetadataQuery metadataQuery
        ) {
            List<Double> widths = new ArrayList<>(scan.getRowType().getFieldCount());
            boolean known = false;
            for (int index = 0; index < scan.getRowType().getFieldCount(); index++) {
                FederationColumnStatistics column = columnStatistics(scan, index);
                Double width = column == null || column.averageWidthBytes() <= 0
                    ? null
                    : (double) column.averageWidthBytes();
                widths.add(width);
                known |= width != null;
            }
            return known ? Collections.unmodifiableList(widths) : null;
        }
    }
}
