"""DSL 表达式层 - 纯 Python 实现，无 pandas/polars 依赖。

提供因子表达式的符号化表示能力，通过重载运算符实现
用户端无感知的公式编写。
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Any, List, Union


class Node(ABC):
    """表达式节点基类。

    所有因子表达式组件的抽象基类，提供运算符重载能力。
    子类需要实现 __repr__ 方法用于表达式可视化。
    """

    # ==================== 算术运算符重载 ====================

    def __add__(self, other: Any) -> BinaryOpNode:
        """加法: self + other"""
        return BinaryOpNode("+", self, _ensure_node(other))

    def __radd__(self, other: Any) -> BinaryOpNode:
        """右加法: other + self"""
        return BinaryOpNode("+", _ensure_node(other), self)

    def __sub__(self, other: Any) -> BinaryOpNode:
        """减法: self - other"""
        return BinaryOpNode("-", self, _ensure_node(other))

    def __rsub__(self, other: Any) -> BinaryOpNode:
        """右减法: other - self"""
        return BinaryOpNode("-", _ensure_node(other), self)

    def __mul__(self, other: Any) -> BinaryOpNode:
        """乘法: self * other"""
        return BinaryOpNode("*", self, _ensure_node(other))

    def __rmul__(self, other: Any) -> BinaryOpNode:
        """右乘法: other * self"""
        return BinaryOpNode("*", _ensure_node(other), self)

    def __truediv__(self, other: Any) -> BinaryOpNode:
        """除法: self / other"""
        return BinaryOpNode("/", self, _ensure_node(other))

    def __rtruediv__(self, other: Any) -> BinaryOpNode:
        """右除法: other / self"""
        return BinaryOpNode("/", _ensure_node(other), self)

    def __pow__(self, other: Any) -> BinaryOpNode:
        """幂运算: self ** other"""
        return BinaryOpNode("**", self, _ensure_node(other))

    def __rpow__(self, other: Any) -> BinaryOpNode:
        """右幂运算: other ** self"""
        return BinaryOpNode("**", _ensure_node(other), self)

    def __floordiv__(self, other: Any) -> BinaryOpNode:
        """整除: self // other"""
        return BinaryOpNode("//", self, _ensure_node(other))

    def __rfloordiv__(self, other: Any) -> BinaryOpNode:
        """右整除: other // self"""
        return BinaryOpNode("//", _ensure_node(other), self)

    def __mod__(self, other: Any) -> BinaryOpNode:
        """取模: self % other"""
        return BinaryOpNode("%", self, _ensure_node(other))

    def __rmod__(self, other: Any) -> BinaryOpNode:
        """右取模: other % self"""
        return BinaryOpNode("%", _ensure_node(other), self)

    # ==================== 一元运算符重载 ====================

    def __neg__(self) -> UnaryOpNode:
        """取负: -self"""
        return UnaryOpNode("-", self)

    def __pos__(self) -> UnaryOpNode:
        """取正: +self"""
        return UnaryOpNode("+", self)

    def __abs__(self) -> UnaryOpNode:
        """绝对值: abs(self)"""
        return UnaryOpNode("abs", self)

    # ==================== 比较运算符重载 ====================

    def __eq__(self, other: Any) -> BinaryOpNode:
        """等于: self == other"""
        return BinaryOpNode("==", self, _ensure_node(other))

    def __ne__(self, other: Any) -> BinaryOpNode:
        """不等于: self != other"""
        return BinaryOpNode("!=", self, _ensure_node(other))

    def __lt__(self, other: Any) -> BinaryOpNode:
        """小于: self < other"""
        return BinaryOpNode("<", self, _ensure_node(other))

    def __le__(self, other: Any) -> BinaryOpNode:
        """小于等于: self <= other"""
        return BinaryOpNode("<=", self, _ensure_node(other))

    def __gt__(self, other: Any) -> BinaryOpNode:
        """大于: self > other"""
        return BinaryOpNode(">", self, _ensure_node(other))

    def __ge__(self, other: Any) -> BinaryOpNode:
        """大于等于: self >= other"""
        return BinaryOpNode(">=", self, _ensure_node(other))

    # ==================== 位运算符重载（用于逻辑运算） ====================

    def __and__(self, other: Any) -> BinaryOpNode:
        """逻辑与: self & other"""
        return BinaryOpNode("&", self, _ensure_node(other))

    def __rand__(self, other: Any) -> BinaryOpNode:
        """右逻辑与: other & self"""
        return BinaryOpNode("&", _ensure_node(other), self)

    def __or__(self, other: Any) -> BinaryOpNode:
        """逻辑或: self | other"""
        return BinaryOpNode("|", self, _ensure_node(other))

    def __ror__(self, other: Any) -> BinaryOpNode:
        """右逻辑或: other | self"""
        return BinaryOpNode("|", _ensure_node(other), self)

    # ==================== 抽象方法 ====================

    @abstractmethod
    def __repr__(self) -> str:
        """返回表达式的字符串表示。"""
        pass


class Symbol(Node):
    """符号节点，代表一个命名变量（如 close, open 等）。

    Attributes:
        name: 符号名称，用于标识该变量
    """

    def __init__(self, name: str) -> None:
        """初始化符号节点。

        Args:
            name: 符号名称，如 'close', 'open', 'volume' 等
        """
        self.name = name

    def __repr__(self) -> str:
        """返回符号名称。"""
        return self.name

    def __hash__(self) -> int:
        """支持作为字典键使用。"""
        return hash(self.name)

    def __eq__(self, other: object) -> bool:
        """符号相等性比较。"""
        if not isinstance(other, Symbol):
            return NotImplemented
        return self.name == other.name


class Constant(Node):
    """常量节点，代表一个数值常量。

    Attributes:
        value: 常量数值
    """

    def __init__(self, value: Union[int, float]) -> None:
        """初始化常量节点。

        Args:
            value: 常量数值
        """
        self.value = value

    def __repr__(self) -> str:
        """返回常量值的字符串表示。"""
        return str(self.value)


class BinaryOpNode(Node):
    """二元运算节点，表示两个操作数之间的运算。

    Attributes:
        op: 运算符，如 '+', '-', '*', '/' 等
        left: 左操作数
        right: 右操作数
    """

    def __init__(self, op: str, left: Node, right: Node) -> None:
        """初始化二元运算节点。

        Args:
            op: 运算符字符串
            left: 左操作数节点
            right: 右操作数节点
        """
        self.op = op
        self.left = left
        self.right = right

    def __repr__(self) -> str:
        """返回带括号的二元运算表达式。"""
        return f"({self.left} {self.op} {self.right})"


class UnaryOpNode(Node):
    """一元运算节点，表示对单个操作数的运算。

    Attributes:
        op: 运算符，如 '-', '+', 'abs' 等
        operand: 操作数
    """

    def __init__(self, op: str, operand: Node) -> None:
        """初始化一元运算节点。

        Args:
            op: 运算符字符串
            operand: 操作数节点
        """
        self.op = op
        self.operand = operand

    def __repr__(self) -> str:
        """返回一元运算表达式。"""
        if self.op in ("+", "-"):
            return f"({self.op}{self.operand})"
        return f"{self.op}({self.operand})"


class FunctionNode(Node):
    """函数调用节点，表示一个函数调用。

    Attributes:
        func_name: 函数名称
        args: 函数参数列表
    """

    def __init__(self, func_name: str, *args: Any) -> None:
        """初始化函数调用节点。

        Args:
            func_name: 函数名称，如 'ts_mean', 'cs_rank' 等
            *args: 函数参数，可以是 Node 或其他类型
        """
        self.func_name = func_name
        # 将所有参数转换为节点类型
        self.args: List[Node] = [_ensure_node(arg) for arg in args]

    def __repr__(self) -> str:
        """返回函数调用表达式。"""
        args_str = ", ".join(repr(arg) for arg in self.args)
        return f"{self.func_name}({args_str})"


# ==================== 辅助函数 ====================


def _ensure_node(value: Any) -> Node:
    """确保值是一个 Node 节点。

    如果值已经是 Node 类型，直接返回；
    如果是数值类型，包装为 Constant 节点；
    如果是字符串类型，包装为 Symbol 节点；
    否则抛出类型错误。

    Args:
        value: 任意值

    Returns:
        Node: 对应的节点对象

    Raises:
        TypeError: 当值无法转换为节点时
    """
    if isinstance(value, Node):
        return value
    if isinstance(value, (int, float)):
        return Constant(value)
    if isinstance(value, str):
        return Symbol(value)
    raise TypeError(f"无法将类型 {type(value).__name__} 转换为 Node")
