"""AST 优化器 - 表达式拍平。

本模块实现将嵌套的窗口函数表达式自动提取为中间临时因子，
解决多维窗口函数（over）嵌套导致计算为空的问题。

核心思想：
    通过 AST 变换，将嵌套在窗口函数内的窗口函数表达式提取出来，
    作为独立的临时因子先行计算，然后主表达式引用这些临时因子。

示例：
    原始表达式: cs_rank(ts_delay(close, 1))
    拍平后:
        - 临时因子: __tmp_0 = ts_delay(close, 1)
        - 主表达式: cs_rank(__tmp_0)
"""

import threading
from typing import Dict, Tuple

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


# 模块级全局计数器，用于生成唯一的临时因子名称
_global_counter: int = 0
_counter_lock = threading.Lock()


def _get_next_counter() -> int:
    """获取下一个全局计数器值。

    Returns:
        递增后的全局计数器值
    """
    global _global_counter
    with _counter_lock:
        _global_counter += 1
        return _global_counter


class ExpressionFlattener:
    """表达式拍平器。

    遍历 AST 并自动提取嵌套的窗口函数为独立临时因子。

    Attributes:
        _extracted_nodes: 存储已提取的临时因子字典
    """

    def __init__(self) -> None:
        """初始化拍平器。"""
        self._extracted_nodes: Dict[str, Node] = {}

    def _generate_temp_name(self) -> str:
        """生成唯一的临时因子名称。

        使用模块级全局计数器确保跨因子注册时的唯一性。

        Returns:
            格式为 "__tmp_X" 的临时名称，其中 X 是全局递增数字
        """
        return f"__tmp_{_get_next_counter()}"

    def _is_window_function(self, func_name: str) -> bool:
        """判断是否为窗口函数。

        窗口函数以 "ts_"（时序）或 "cs_"（截面）开头。

        Args:
            func_name: 函数名称

        Returns:
            是否是窗口函数
        """
        return func_name.startswith("ts_") or func_name.startswith("cs_")

    def flatten(self, node: Node) -> Tuple[Node, Dict[str, Node]]:
        """拍平表达式。

        遍历 AST，将嵌套的窗口函数提取为临时因子。

        Args:
            node: 原始表达式根节点

        Returns:
            Tuple[拍平后的主表达式节点, 临时因子字典]
            临时因子字典: {临时名称 -> 被提取的节点}

        Example:
            >>> flattener = ExpressionFlattener()
            >>> from src.factors.dsl import Symbol, FunctionNode
            >>> close = Symbol("close")
            >>> expr = FunctionNode("cs_rank", FunctionNode("ts_delay", close, 1))
            >>> flat_expr, tmp_factors = flattener.flatten(expr)
            >>> # flat_expr = cs_rank(__tmp_0)
            >>> # tmp_factors = {"__tmp_0": ts_delay(close, 1)}
        """
        # 重置状态（只重置提取的节点字典，计数器保持全局递增）
        self._extracted_nodes = {}

        # 从根节点开始遍历，初始状态为不在窗口函数内部
        new_node = self._flatten_recursive(node, inside_window=False)

        return new_node, self._extracted_nodes.copy()

    def _flatten_recursive(self, node: Node, inside_window: bool) -> Node:
        """递归拍平节点。

        Args:
            node: 当前处理的节点
            inside_window: 当前是否处于窗口函数内部

        Returns:
            处理后的节点（可能是原节点或替换为 Symbol）
        """
        # Symbol 和 Constant 是叶子节点，直接返回
        if isinstance(node, Symbol):
            return node

        if isinstance(node, Constant):
            return node

        # 处理二元运算节点
        if isinstance(node, BinaryOpNode):
            return self._flatten_binary_op(node, inside_window)

        # 处理一元运算节点
        if isinstance(node, UnaryOpNode):
            return self._flatten_unary_op(node, inside_window)

        # 处理函数调用节点
        if isinstance(node, FunctionNode):
            return self._flatten_function(node, inside_window)

        # 未知节点类型，直接返回
        return node

    def _flatten_binary_op(self, node: BinaryOpNode, inside_window: bool) -> Node:
        """拍平二元运算节点。

        Args:
            node: 二元运算节点
            inside_window: 当前是否处于窗口函数内部

        Returns:
            处理后的节点
        """
        # 递归处理左右子节点
        new_left = self._flatten_recursive(node.left, inside_window)
        new_right = self._flatten_recursive(node.right, inside_window)

        # 如果子节点没有变化，返回原节点
        if new_left is node.left and new_right is node.right:
            return node

        # 创建新的二元运算节点
        return BinaryOpNode(node.op, new_left, new_right)

    def _flatten_unary_op(self, node: UnaryOpNode, inside_window: bool) -> Node:
        """拍平一元运算节点。

        Args:
            node: 一元运算节点
            inside_window: 当前是否处于窗口函数内部

        Returns:
            处理后的节点
        """
        # 递归处理操作数
        new_operand = self._flatten_recursive(node.operand, inside_window)

        # 如果操作数没有变化，返回原节点
        if new_operand is node.operand:
            return node

        # 创建新的一元运算节点
        return UnaryOpNode(node.op, new_operand)

    def _flatten_function(self, node: FunctionNode, inside_window: bool) -> Node:
        """拍平函数调用节点。

        修正为后序遍历（Bottom-Up）：先递归拍平参数，再决定是否提取当前节点。
        确保深层嵌套（如 3层以上）也能被彻底逐层拆解。

        Args:
            node: 函数调用节点
            inside_window: 当前是否处于窗口函数内部

        Returns:
            处理后的节点
        """
        is_window = self._is_window_function(node.func_name)
        next_inside_window = inside_window or is_window

        # 1. 优先递归处理所有参数
        new_args = []
        has_change = False
        for arg in node.args:
            new_arg = self._flatten_recursive(arg, next_inside_window)
            new_args.append(new_arg)
            if new_arg is not arg:
                has_change = True

        # 2. 只有当参数发生变化时，才创建新的当前节点
        current_node = FunctionNode(node.func_name, *new_args) if has_change else node

        # 3. 判断是否需要提取（此时子节点肯定已经被彻底拍平了）
        if inside_window and is_window:
            temp_name = self._generate_temp_name()
            self._extracted_nodes[temp_name] = current_node
            return Symbol(temp_name)

        return current_node


def flatten_expression(node: Node) -> Tuple[Node, Dict[str, Node]]:
    """便捷函数 - 拍平表达式。

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

    Returns:
        Tuple[拍平后的主表达式节点, 临时因子字典]

    Example:
        >>> from src.factors.dsl import Symbol, FunctionNode
        >>> close = Symbol("close")
        >>> expr = FunctionNode("cs_rank", FunctionNode("ts_delay", close, 1))
        >>> flat_expr, tmp_factors = flatten_expression(expr)
    """
    flattener = ExpressionFlattener()
    return flattener.flatten(node)
