"""因子管理器

管理多种来源的因子：
- metadata 中注册的因子
- DSL 表达式定义的因子
- Label 因子
- 排除的因子列表
"""

from typing import Dict, List, Optional, Any
import polars as pl

from src.factors import FactorEngine


class FactorManager:
    """因子管理器

    统一管理多种来源的因子注册和准备：
    1. metadata 中已注册的因子（通过名称引用）
    2. DSL 表达式定义的因子（动态注册）
    3. Label 因子（通过表达式定义）
    4. 排除的因子列表（从最终列表中移除）

    Attributes:
        selected_factors: 从 metadata 中选择的因子名称列表
        factor_definitions: DSL 表达式定义的因子字典 {name: dsl_expression}
        label_factor: Label 因子定义 {name: dsl_expression}
        excluded_factors: 需要排除的因子名称列表
        registered_factors: 已注册到 FactorEngine 的因子列表
    """

    def __init__(
        self,
        selected_factors: List[str],
        factor_definitions: Dict[str, str],
        label_factor: Dict[str, str],
        excluded_factors: Optional[List[str]] = None,
    ):
        """初始化因子管理器

        Args:
            selected_factors: 从 metadata 中选择的因子名称列表
            factor_definitions: DSL 表达式定义的因子字典
            label_factor: Label 因子定义字典
            excluded_factors: 需要排除的因子名称列表
        """
        self.selected_factors = selected_factors or []
        self.factor_definitions = factor_definitions or {}
        self.label_factor = label_factor or {}
        self.excluded_factors = excluded_factors or []
        self.registered_factors: List[str] = []

    def register_to_engine(
        self,
        engine: FactorEngine,
        verbose: bool = True,
    ) -> List[str]:
        """注册所有因子到 FactorEngine

        按以下顺序注册：
        1. metadata 中的因子（通过名称从 metadata 加载）
        2. DSL 表达式定义的因子（使用 add_factor 注册）
        3. Label 因子（使用 add_factor 注册）
        4. 排除指定的因子

        Args:
            engine: FactorEngine 实例
            verbose: 是否打印注册信息

        Returns:
            最终的特征列名列表（已排除指定因子）
        """
        if verbose:
            print("\n" + "=" * 80)
            print("因子注册")
            print("=" * 80)

        # Step 1: 从 metadata 注册选中的因子
        if verbose:
            print(f"\n[1/4] 从 metadata 注册 {len(self.selected_factors)} 个因子...")

        feature_cols = []
        for factor_name in self.selected_factors:
            try:
                engine.add_factor(factor_name)
                feature_cols.append(factor_name)
                if verbose:
                    print(f"  [OK] {factor_name}")
            except Exception as e:
                if verbose:
                    print(f"  [FAIL] {factor_name}: {e}")

        # Step 2: 注册 DSL 定义的因子
        if self.factor_definitions:
            if verbose:
                print(f"\n[2/4] 注册 {len(self.factor_definitions)} 个 DSL 定义因子...")

            for factor_name, dsl_expr in self.factor_definitions.items():
                if factor_name not in self.excluded_factors:
                    try:
                        engine.add_factor(factor_name, dsl_expr)
                        feature_cols.append(factor_name)
                        if verbose:
                            print(f"  ✓ {factor_name}: {dsl_expr[:50]}...")
                    except Exception as e:
                        if verbose:
                            print(f"  ✗ {factor_name}: {e}")

        # Step 3: 注册 Label 因子
        if self.label_factor:
            if verbose:
                print(f"\n[3/4] 注册 Label 因子...")

            for factor_name, dsl_expr in self.label_factor.items():
                try:
                    engine.add_factor(factor_name, dsl_expr)
                    if verbose:
                        print(f"  ✓ Label: {factor_name}")
                except Exception as e:
                    if verbose:
                        print(f"  ✗ Label {factor_name}: {e}")

        # Step 4: 排除指定因子
        if self.excluded_factors:
            if verbose:
                print(f"\n[4/4] 排除 {len(self.excluded_factors)} 个因子...")

            original_count = len(feature_cols)
            feature_cols = [f for f in feature_cols if f not in self.excluded_factors]
            excluded_count = original_count - len(feature_cols)

            if verbose:
                print(f"  排除 {excluded_count} 个因子")
                for f in self.excluded_factors:
                    if f in self.selected_factors or f in self.factor_definitions:
                        print(f"    - {f}")

        self.registered_factors = feature_cols

        if verbose:
            print(f"\n[结果] 最终特征数: {len(feature_cols)}")
            print("=" * 80)

        return feature_cols

    def get_feature_cols(self) -> List[str]:
        """获取已注册的特征列名列表

        Returns:
            特征列名列表
        """
        return self.registered_factors

    def get_label_col(self) -> Optional[str]:
        """获取 Label 列名

        Returns:
            Label 列名，如果没有则返回 None
        """
        if self.label_factor:
            return list(self.label_factor.keys())[0]
        return None
