package com.easyagents.federation.sql.runtime;

import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import org.apache.calcite.DataContext;
import org.apache.calcite.linq4j.Enumerable;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeFactory;
import org.apache.calcite.schema.ScannableTable;
import org.apache.calcite.schema.Statistic;
import org.apache.calcite.schema.Statistics;
import org.apache.calcite.schema.impl.AbstractTable;

/**
 * Calcite 本地计划中的无状态分片输入表。
 */
public final class FederationFragmentTable extends AbstractTable implements ScannableTable {

    private final String fragmentId;
    private final RelDataType rowType;
    private final double estimatedRows;

    /**
     * 创建分片输入表。
     *
     * @param fragmentId 分片标识
     * @param rowType Calcite 行类型
     * @param estimatedRows 分片输出估算行数
     */
    FederationFragmentTable(String fragmentId, RelDataType rowType, double estimatedRows) {
        this.fragmentId = fragmentId;
        this.rowType = rowType;
        this.estimatedRows = Math.max(0D, estimatedRows);
    }

    /**
     * 创建缺少成本统计的兼容分片表。
     *
     * @param fragmentId 分片标识
     * @param rowType Calcite 行类型
     */
    FederationFragmentTable(String fragmentId, RelDataType rowType) {
        this(fragmentId, rowType, 100D);
    }

    /** {@inheritDoc} */
    @Override
    public RelDataType getRowType(RelDataTypeFactory typeFactory) {
        return rowType;
    }

    /** {@inheritDoc} */
    @Override
    public Statistic getStatistic() {
        return Statistics.of(estimatedRows, java.util.List.of());
    }

    /** {@inheritDoc} */
    @Override
    public Enumerable<Object[]> scan(DataContext root) {
        Object value = root.get(FederationExecutionSession.DATA_CONTEXT_KEY);
        if (!(value instanceof FederationExecutionSession session)) {
            throw new FederationSqlException(
                FederationSqlErrorCode.EXECUTION_FAILED,
                "federation execution session is missing from Calcite DataContext"
            );
        }
        return session.scan(fragmentId);
    }
}
