package com.easyagents.federation.sql.runtime;

import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import com.easyagents.federation.sql.federation.FederationLogicalTableDefinition;
import com.easyagents.federation.sql.federation.FederationQueryScopeDefinition;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import org.apache.calcite.sql.SqlCall;
import org.apache.calcite.sql.SqlIdentifier;
import org.apache.calcite.sql.SqlJoin;
import org.apache.calcite.sql.SqlKind;
import org.apache.calcite.sql.SqlNode;
import org.apache.calcite.sql.SqlSelect;
import org.apache.calcite.sql.SqlWith;
import org.apache.calcite.sql.SqlWithItem;
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
import org.apache.calcite.sql.util.SqlShuttle;

/**
 * 使用 Calcite AST 将查询范围逻辑表名解析为物理 Binding 表路径。
 */
final class LogicalTableSqlResolver {

    private final Map<String, FederationLogicalTableDefinition> tablesByName;
    private final Map<String, FederationLogicalTableDefinition> tablesByPath;

    private LogicalTableSqlResolver(List<FederationLogicalTableDefinition> tables) {
        Map<String, FederationLogicalTableDefinition> byName = new HashMap<>();
        Map<String, FederationLogicalTableDefinition> byPath = new HashMap<>();
        for (FederationLogicalTableDefinition table : tables) {
            byName.put(normalize(table.logicalName()), table);
            byPath.put(pathKey(
                table.bindingName(),
                table.schemaName(),
                table.logicalName()
            ), table);
        }
        this.tablesByName = Map.copyOf(byName);
        this.tablesByPath = Map.copyOf(byPath);
    }

    /**
     * 解析查询范围中声明的逻辑表；未声明映射时原样返回。
     *
     * @param parsed Calcite 已解析 SQL AST
     * @param scope 查询范围
     * @return 可交给 Calcite Validator 的物理表 AST
     */
    static SqlNode resolve(SqlNode parsed, FederationQueryScopeDefinition scope) {
        if (scope.logicalTables().isEmpty()) {
            return parsed;
        }
        LogicalTableSqlResolver resolver = new LogicalTableSqlResolver(scope.logicalTables());
        return parsed.accept(resolver.new ResolverShuttle(Set.of()));
    }

    /**
     * 返回短逻辑表实际绑定的物理数据源。
     *
     * @param logicalName 逻辑表名
     * @param scope 查询范围
     * @return Binding 名称；不存在时返回 {@code null}
     */
    static String bindingForShortName(
        String logicalName,
        FederationQueryScopeDefinition scope
    ) {
        if (scope.logicalTables().isEmpty()) {
            return null;
        }
        return scope.logicalTables().stream()
            .filter(table -> table.logicalName().equalsIgnoreCase(logicalName))
            .map(FederationLogicalTableDefinition::bindingName)
            .findFirst()
            .orElse(null);
    }

    /**
     * 创建大小写不敏感且无分隔歧义的三段路径键。
     *
     * @param binding Binding 名称
     * @param schema Schema 名称
     * @param table 逻辑表名
     * @return 路径键
     */
    private static String pathKey(String binding, String schema, String table) {
        return normalize(binding) + '\u0000' + normalize(schema) + '\u0000' + normalize(table);
    }

    /**
     * 按 Calcite 未加引号标识符规则归一化名称。
     *
     * @param value 原始名称
     * @return 大写名称
     */
    private static String normalize(String value) {
        return value.toUpperCase(Locale.ROOT);
    }

    /**
     * 创建未知逻辑表校验异常。
     *
     * @param identifier 未命中的表标识符
     * @return SQL 校验异常
     */
    private FederationSqlException unknownTable(SqlIdentifier identifier) {
        return new FederationSqlException(
            FederationSqlErrorCode.SQL_VALIDATION_FAILED,
            "logical table is not declared in the query scope: "
                + String.join(".", identifier.names)
        );
    }

    /**
     * 仅在关系来源位置解析表名，并维护 CTE 词法作用域。
     */
    private final class ResolverShuttle extends SqlShuttle {

        private final Set<String> visibleCtes;

        private ResolverShuttle(Set<String> visibleCtes) {
            this.visibleCtes = visibleCtes;
        }

        /** {@inheritDoc} */
        @Override
        public SqlNode visit(SqlCall call) {
            if (call instanceof SqlWith with) {
                return rewriteWith(with);
            }
            if (call instanceof SqlSelect) {
                SqlSelect rewritten = (SqlSelect) super.visit(call);
                rewritten.setFrom(rewriteFrom(rewritten.getFrom(), true));
                return rewritten;
            }
            return super.visit(call);
        }

        /**
         * 将四段逻辑列限定名转换为隐式逻辑表别名限定名。
         *
         * @param identifier SQL 标识符
         * @return 原标识符或两段逻辑列限定名
         */
        @Override
        public SqlNode visit(SqlIdentifier identifier) {
            if (identifier.names.size() != 4) {
                return identifier;
            }
            FederationLogicalTableDefinition table = tablesByPath.get(pathKey(
                identifier.names.get(0),
                identifier.names.get(1),
                identifier.names.get(2)
            ));
            if (table == null) {
                return identifier;
            }
            return new SqlIdentifier(
                List.of(normalize(table.logicalName()), identifier.names.get(3)),
                identifier.getParserPosition()
            );
        }

        /**
         * 按声明顺序重写 CTE，并维护其词法可见范围。
         *
         * @param with CTE 节点
         * @return 重写后的 CTE
         */
        private SqlNode rewriteWith(SqlWith with) {
            Set<String> withScope = new HashSet<>(visibleCtes);
            for (SqlNode node : with.withList) {
                SqlWithItem item = (SqlWithItem) node;
                Set<String> itemScope = new HashSet<>(withScope);
                if (item.recursive != null && item.recursive.booleanValue()) {
                    itemScope.add(normalize(item.name.getSimple()));
                }
                item.query = item.query.accept(new ResolverShuttle(itemScope));
                withScope.add(normalize(item.name.getSimple()));
            }
            with.body = with.body.accept(new ResolverShuttle(withScope));
            return with;
        }

        /**
         * 递归重写 FROM 关系来源中的逻辑表。
         *
         * @param from 关系来源
         * @param addImplicitAlias 是否补充逻辑表隐式别名
         * @return 重写后的关系来源
         */
        private SqlNode rewriteFrom(SqlNode from, boolean addImplicitAlias) {
            if (from == null) {
                return null;
            }
            if (from instanceof SqlIdentifier identifier) {
                return resolveIdentifier(identifier, addImplicitAlias);
            }
            if (from instanceof SqlJoin join) {
                join.setLeft(rewriteFrom(join.getLeft(), true));
                join.setRight(rewriteFrom(join.getRight(), true));
                return join;
            }
            if (from instanceof SqlSelect || from instanceof SqlWith) {
                return from.accept(new ResolverShuttle(visibleCtes));
            }
            if (from instanceof SqlCall call && isRelationWrapper(call.getKind())) {
                if (call.operandCount() > 0) {
                    boolean wrappedByExplicitAlias = call.getKind() == SqlKind.AS;
                    call.setOperand(0, rewriteFrom(call.operand(0), !wrappedByExplicitAlias));
                }
                return call;
            }
            return from;
        }

        /**
         * 将短名或三段逻辑表名解析为物理 Binding 路径。
         *
         * @param identifier 逻辑表标识符
         * @param addImplicitAlias 是否补充逻辑表隐式别名
         * @return 物理表路径或带别名的关系节点
         * @throws FederationSqlException 逻辑表未声明时抛出
         */
        private SqlNode resolveIdentifier(
            SqlIdentifier identifier,
            boolean addImplicitAlias
        ) {
            FederationLogicalTableDefinition table;
            if (identifier.isSimple()) {
                String name = identifier.getSimple();
                if (visibleCtes.contains(normalize(name))) {
                    return identifier;
                }
                table = tablesByName.get(normalize(name));
            } else if (identifier.names.size() == 3) {
                table = tablesByPath.get(pathKey(
                    identifier.names.get(0),
                    identifier.names.get(1),
                    identifier.names.get(2)
                ));
            } else {
                throw unknownTable(identifier);
            }
            if (table == null) {
                throw unknownTable(identifier);
            }
            SqlIdentifier physical = new SqlIdentifier(
                List.of(table.bindingName(), table.schemaName(), table.sourceTableName()),
                identifier.getParserPosition()
            );
            if (!addImplicitAlias) {
                return physical;
            }
            return SqlStdOperatorTable.AS.createCall(
                identifier.getParserPosition(),
                physical,
                new SqlIdentifier(
                    normalize(table.logicalName()),
                    identifier.getParserPosition()
                )
            );
        }

        /**
         * 判断调用节点是否仅包装一个关系来源。
         *
         * @param kind Calcite 节点类型
         * @return 是否为关系包装节点
         */
        private boolean isRelationWrapper(SqlKind kind) {
            return kind == SqlKind.AS
                || kind == SqlKind.LATERAL
                || kind == SqlKind.TABLESAMPLE
                || kind == SqlKind.SNAPSHOT;
        }
    }
}
