"""
历史数据预切片优化：通过 AST 静态分析提取 DSL 表达式中的最大时间窗口，
仅保留计算所需的最近 N 个交易日数据，减少时序因子的滚动计算量。
"""

import ast
from typing import Optional

import polars as pl
import ldfunctions


def _build_window_index_map() -> dict[str, tuple[tuple[int, ...], str]]:
    """从 ldfunctions 中收集所有标注了 @ts_window 的函数及其窗口参数元数据

    Returns:
        {函数名: (窗口参数位置索引元组, 聚合方式)}
    """
    return {
        name: (obj._ts_window_indices, getattr(obj, '_ts_window_combine', 'max'))
        for name, obj in ldfunctions.__dict__.items()
        if callable(obj) and hasattr(obj, '_ts_window_indices')
    }


# 模块级缓存，避免每次调用重建
_WINDOW_INDEX_MAP: dict[str, tuple[tuple[int, ...], str]] | None = None


def _get_window_index_map() -> dict[str, tuple[tuple[int, ...], str]]:
    global _WINDOW_INDEX_MAP
    if _WINDOW_INDEX_MAP is None:
        _WINDOW_INDEX_MAP = _build_window_index_map()
    return _WINDOW_INDEX_MAP


def _extract_call_window(node: ast.Call, window_map: dict) -> int:
    """从单个 AST 函数调用节点中提取有效窗口大小"""
    func_name = node.func.id if isinstance(node.func, ast.Name) else None
    if func_name not in window_map:
        return 0

    indices, combine = window_map[func_name]
    windows = []
    for idx in indices:
        if idx < len(node.args):
            arg = node.args[idx]
            if isinstance(arg, ast.Constant) and isinstance(arg.value, (int, float)):
                windows.append(int(arg.value))

    if not windows:
        return 0
    return sum(windows) if combine == 'sum' else max(windows)


def extract_max_window(computed_factors: Optional[list[list[dict]]]) -> int:
    """从自定义因子 DSL 表达式中提取最大时间窗口

    各 stage 的窗口取 max（同 stage 内因子并行计算，互不依赖）。
    各 stage 之间的窗口求 sum（后续 stage 可能依赖前序 stage 的结果，窗口叠加）。

    Args:
        computed_factors: 分层的自定义因子列表，如 [[stage1_factors], [stage2_factors]]

    Returns:
        需要的最大历史交易日数
    """
    if not computed_factors:
        return 0

    window_map = _get_window_index_map()
    stage_windows = []

    for stage in computed_factors:
        stage_max = 0
        for factor in stage:
            expr_str = factor.get('expr', '')
            try:
                tree = ast.parse(expr_str, mode='eval')
            except SyntaxError:
                continue

            for node in ast.walk(tree):
                if isinstance(node, ast.Call):
                    w = _extract_call_window(node, window_map)
                    stage_max = max(stage_max, w)

        stage_windows.append(stage_max)

    return sum(stage_windows)


def slice_by_fresh_index(
    lf: pl.LazyFrame,
    max_window: int,
    buffer: int = 30,
) -> pl.LazyFrame:
    """按 _fresh_index 列预切片历史数据（零开销，predicate pushdown 友好）

    _fresh_index 表示距最新交易日的天数（1=最新，1754=最早），
    直接作为 parquet 过滤条件，polars 可跳过不需要的 row groups。

    Args:
        lf: 历史数据 LazyFrame（需包含 _fresh_index 列）
        max_window: 需要的最大历史交易日数
        buffer: 额外安全余量天数
    """
    return lf.filter(pl.col('_fresh_index') <= max_window + buffer)
