"""AST 编译器模块 - 提供依赖提取和代码生成功能。

本模块实现 AST 遍历器模式，用于从 DSL 表达式中提取依赖的符号。
"""

from typing import Set, Optional

from src.factors.dsl import Node, Symbol, BinaryOpNode, UnaryOpNode, FunctionNode


class DependencyExtractor:
    """依赖提取器 - 使用访问者模式遍历 AST 节点。

    递归遍历表达式树，提取所有 Symbol 节点的名称。
    支持 BinaryOpNode、UnaryOpNode 和 FunctionNode 的递归遍历。

    Example:
        >>> from src.factors.dsl import Symbol, FunctionNode
        >>> close = Symbol("close")
        >>> pe_ratio = Symbol("pe_ratio")
        >>> alpha = FunctionNode("cs_rank", close / pe_ratio)
        >>> deps = DependencyExtractor.extract_dependencies(alpha)
        >>> print(deps)
        {'close', 'pe_ratio'}
    """

    def __init__(self, ignore_symbols: Optional[Set[str]] = None) -> None:
        """初始化依赖提取器。

        Args:
            ignore_symbols: 需要忽略的符号集合（如已注册的因子名）
        """
        self.dependencies: Set[str] = set()
        self.ignore_symbols: Set[str] = ignore_symbols or set()

    def visit(self, node: Node) -> None:
        """访问节点，根据节点类型分发到具体处理方法。

        Args:
            node: AST 节点
        """
        if isinstance(node, Symbol):
            self._visit_symbol(node)
        elif isinstance(node, BinaryOpNode):
            self._visit_binary_op(node)
        elif isinstance(node, UnaryOpNode):
            self._visit_unary_op(node)
        elif isinstance(node, FunctionNode):
            self._visit_function(node)
        # Constant 节点不包含依赖，无需处理

    def _visit_symbol(self, node: Symbol) -> None:
        """访问 Symbol 节点，提取符号名称。

        排除临时因子（以 __tmp_ 开头的符号）和已在免疫名单中的因子。

        Args:
            node: 符号节点
        """
        # 排除临时因子引用 和 已在免疫名单中的因子
        if not node.name.startswith("__tmp_") and node.name not in self.ignore_symbols:
            self.dependencies.add(node.name)

    def _visit_binary_op(self, node: BinaryOpNode) -> None:
        """访问 BinaryOpNode 节点，递归遍历左右子节点。

        Args:
            node: 二元运算节点
        """
        self.visit(node.left)
        self.visit(node.right)

    def _visit_unary_op(self, node: UnaryOpNode) -> None:
        """访问 UnaryOpNode 节点，递归遍历操作数。

        Args:
            node: 一元运算节点
        """
        self.visit(node.operand)

    def _visit_function(self, node: FunctionNode) -> None:
        """访问 FunctionNode 节点，递归遍历所有参数。

        Args:
            node: 函数调用节点
        """
        for arg in node.args:
            self.visit(arg)

    def extract(self, node: Node) -> Set[str]:
        """从 AST 节点中提取所有依赖的符号名称。

        Args:
            node: 表达式树的根节点

        Returns:
            依赖的符号名称集合
        """
        self.dependencies.clear()
        self.visit(node)
        return self.dependencies.copy()

    @classmethod
    def extract_dependencies(
        cls, node: Node, ignore_symbols: Optional[Set[str]] = None
    ) -> Set[str]:
        """类方法 - 从 AST 节点中提取所有依赖的符号名称。

        这是一个便捷方法，无需手动实例化 DependencyExtractor。

        Args:
            node: 表达式树的根节点
            ignore_symbols: 需要忽略的符号集合（如已注册的因子名）

        Returns:
            依赖的符号名称集合

        Example:
            >>> from src.factors.dsl import Symbol
            >>> close = Symbol("close")
            >>> open_price = Symbol("open")
            >>> expr = close / open_price
            >>> deps = DependencyExtractor.extract_dependencies(expr)
            >>> print(deps)
            {'close', 'open'}
        """
        extractor = cls(ignore_symbols=ignore_symbols)
        return extractor.extract(node)


def extract_dependencies(
    node: Node, ignore_symbols: Optional[Set[str]] = None
) -> Set[str]:
    """单例方法 - 从 AST 节点中提取所有依赖的符号名称。

    这是 DependencyExtractor.extract_dependencies 的便捷包装函数。

    Args:
        node: 表达式树的根节点
        ignore_symbols: 需要忽略的符号集合（如已注册的因子名）

    Returns:
        依赖的符号名称集合

    Example:
        >>> from src.factors.dsl import Symbol, FunctionNode
        >>> close = Symbol("close")
        >>> pe_ratio = Symbol("pe_ratio")
        >>> alpha = FunctionNode("cs_rank", close / pe_ratio)
        >>> deps = extract_dependencies(alpha)
        >>> print(deps)
        {'close', 'pe_ratio'}
    """
    return DependencyExtractor.extract_dependencies(node, ignore_symbols=ignore_symbols)


if __name__ == "__main__":
    # 测试用例: cs_rank(close / pe_ratio)
    from src.factors.dsl import Symbol, FunctionNode

    # 创建符号
    close = Symbol("close")
    pe_ratio = Symbol("pe_ratio")

    # 构建表达式: cs_rank(close / pe_ratio)
    alpha = FunctionNode("cs_rank", close / pe_ratio)

    # 提取依赖
    dependencies = extract_dependencies(alpha)

    print(f"表达式: {alpha}")
    print(f"提取的依赖: {dependencies}")
    print(f"期望依赖: {{'close', 'pe_ratio'}}")
    print(f"验证结果: {dependencies == {'close', 'pe_ratio'}}")
