from typing import Optional

import polars as pl
import ldfunctions

from ldnode.calculator.schemas import (
    ExprType,
    ComputedFactor,
)


def planning_computed_factors(
    frame: pl.LazyFrame | pl.DataFrame,
    builtin_factors: list[str],
    computed_factors: Optional[list[list[ComputedFactor]]],
    exclude_suspension: bool = True,
    root_cols: Optional[set[str]] = None,
) -> pl.LazyFrame:
    """ 规划自定义因子计算 """
    lf = frame.lazy() if isinstance(frame, pl.DataFrame) else frame
    # 若没有自定义因子，则直接返回
    if not computed_factors:
        return lf
    # 排除停牌记录
    if exclude_suspension:
        # 保存停牌记录，后续恢复
        suspension_lf = lf.filter(pl.col('suspension'))
        lf = lf.filter(~pl.col('suspension'))
    # 生成dsl和polars的locals字典
    dsl_locals_dict = {
        # 添加dsl函数
        **{
            name: obj
            for name, obj in ldfunctions.__dict__.items()
            if not name.startswith('__')
        },
        # 添加内置因子
        **{
            factor: pl.col(factor)
            for factor in builtin_factors
        }
    }
    polars_locals_dict = {
        'col': pl.col, 
        'len': pl.len,
    }
    # 停牌记录中需要添加的列
    suspension_columns = [] 
    # 生成计算计划
    for computed_sequence in computed_factors:
        parallel_exprs = []
        for computed_factor in computed_sequence: 
            match computed_factor['expr_type']:
                case ExprType.DSL:
                    try:
                        expr = eval(
                            computed_factor['expr'],
                            {'__builtins__': None},
                            dsl_locals_dict
                        )
                    except TypeError as e:
                        # __builtins__=None 时，未定义变量的查找会落到 None[name]，
                        # 产生 "NoneType object is not subscriptable" 的 TypeError，
                        # 此处转换为可读的 NameError；其他 TypeError（如函数参数数量错误）则直接透传
                        if "'NoneType' object is not subscriptable" in str(e):
                            raise NameError(
                                f"DSL 表达式中存在未定义的因子或函数: {computed_factor['expr']!r}"
                            ) from e
                        raise
                    # 在dsl_locals_dict中添加此自定义因子的引用
                    dsl_locals_dict.update({computed_factor['factor']: pl.col(computed_factor['factor'])})
                case ExprType.POLARS:
                    expr = eval(
                        computed_factor['expr'], 
                        {'__builtins__': None}, 
                        polars_locals_dict
                    )
                case ExprType.SQL:
                    expr = pl.sql_expr(computed_factor['expr'])
                case _:
                    raise ValueError(f"Invalid expr_type: {computed_factor['expr_type']}")
            # 记录根列
            if root_cols is not None and isinstance(expr, pl.Expr):
                root_cols |= set(expr.meta.root_names())
            # 添加因子列
            if isinstance(expr, pl.Expr):
                parallel_exprs.append(expr.alias(computed_factor['factor']))
            else:
                parallel_exprs.append(pl.lit(expr).alias(computed_factor['factor']))
            # 添加停牌记录中需要添加的列
            suspension_columns.append(pl.lit(None).alias(computed_factor['factor']))
        # 添加因子列
        if parallel_exprs:
            lf = lf.with_columns(parallel_exprs)
    # 还原停牌行
    if exclude_suspension:
        # 为停牌记录lf添加相应的schema
        suspension_lf = suspension_lf.with_columns(suspension_columns)
        lf = pl.concat([lf, suspension_lf]).sort('trade_date')
    return lf
    