from typing import Optional, Type

import polars as pl

from ldnode.calculator.schemas import (
    SecurityPoolType,
    SecurityPoolCustomizeType,
    SecurityPool,
    ExcludeOperator,
    BuiltinScope,
)


def _gen_customized_pool_expr(security_pool: SecurityPool) -> Optional[pl.Expr]:
    """ 生成自定义标的池表达式 """
    match security_pool.get('pool_customize_type'):
        case SecurityPoolCustomizeType.FIXED:
            if securities := security_pool.get('pool_securities'):
                return pl.col('code').is_in(securities)
        case SecurityPoolCustomizeType.DYNAMIC_BY_FACTOR:
            if filter_factors := security_pool.get('pool_filter_factors'):
                filter_exprs = [
                    ExcludeOperator(filter_factor['operator']).compare(filter_factor['factor'], filter_factor['value'])
                    for filter_factor in filter_factors
                    if filter_factor['operator'] in ExcludeOperator
                ]
                return pl.Expr.and_(*filter_exprs)
            

def _gen_pool_expr(security_pool: SecurityPool, BulitinScopes: Type[BuiltinScope]) -> Optional[pl.Expr]:
    """ 生成标的池表达式 """
    match security_pool['pool_type']:
        case SecurityPoolType.BUILTIN:
            if builtin_scope := security_pool.get('pool_builtin_scope'):
                if builtin_scope in BulitinScopes:
                    return BulitinScopes(builtin_scope).process()
        case SecurityPoolType.CUSTOMIZED:
            return _gen_customized_pool_expr(security_pool)


def _gen_pool_exprs(security_pools: list[SecurityPool], BulitinScopes: Type[BuiltinScope]) -> list[pl.Expr]:
    """ 生成标的池表达式列表 """
    return list(
        filter(
            lambda x: x is not None,
            [_gen_pool_expr(security_pool, BulitinScopes) for security_pool in security_pools], 
        )
    )


def planning_security_pools(
    frame: pl.LazyFrame | pl.DataFrame,
    BuiltinScopes: Type[BuiltinScope],
    security_pools: Optional[list[SecurityPool]],
) -> pl.LazyFrame:
    """ 
        规划股票池计算 
        前置依赖: 规划百分比因子计算
    """
    lf = frame.lazy() if isinstance(frame, pl.DataFrame) else frame
    # 兜底：确保 code 列为 String 类型（cb_data.pq 中 code 为 Categorical，
    # BuiltinScope 的 str.starts_with 等操作要求 String）
    if 'code' in lf.collect_schema().names():
        lf = lf.with_columns(pl.col('code').cast(pl.String))
    # 若没有标的池，则直接返回
    if not security_pools:
        return lf.with_columns(pl.lit(True).alias('_isin_pool'))
    # 标的池表达式
    if pool_exprs := _gen_pool_exprs(security_pools, BuiltinScopes):
        pool_expr = pl.Expr.or_(*pool_exprs).fill_null(False)
        # 曾入池的标的保留全部数据
        ever_in_pool_codes = (
            lf.filter(pool_expr)
            .select('code')
            .unique()
            .collect()
            .get_column('code')
        )
        # 只保留曾经入选标的池的标的
        return lf.filter(
            pl.col('code').is_in(ever_in_pool_codes)
        ).with_columns(
            (pool_expr).alias('_isin_pool')
        )
    else:
        return lf.with_columns(pl.lit(True).alias('_isin_pool'))
    