from typing import Any
from datetime import datetime

import polars as pl

from ldnode.loader.files import STK_DATA, STK_QUOTATION
from ldnode.calculator.schemas import STK_BUILTIN_FACTORS
from ldnode.calculator.operators import stk_screen
from ldnode.calculator.plans.polars import maybe_resolve_joins
from ldnode.calculator.plans.polars.window_slicer import extract_max_window, slice_by_fresh_index
from .utils import wrap_handler, response
from .models import STKScreenParameters, TaskResponse, NullScoreType


@wrap_handler("选股任务")
def handle_stk_screen(parameters: STKScreenParameters) -> TaskResponse[dict[str, list[dict[str, Any]]]]:
    """ 响应选股任务 """
    # 历史数据预切片：提取所有策略的最大时间窗口，在 scan_parquet 后立即过滤，
    # 让 polars predicate pushdown 跳过不需要的 parquet row groups
    max_window = max(
        (extract_max_window(cfg.get("computed_factors"))
         for cfg in parameters['strategy_configs'].values()),
        default=0,
    )
    lf_history = STK_DATA.output_lazyframe()
    if max_window > 0:
        lf_history = slice_by_fresh_index(
            lf_history, max_window,
            target_date=datetime.strptime(parameters['trade_date'], '%Y-%m-%d').date() if parameters.get('trade_date') else None,
        )
    lf_history = lf_history.filter(~pl.col('suspension')).select(STK_BUILTIN_FACTORS)
    # 处理选股截面
    if parameters.get('trade_date'):
        trade_date = datetime.strptime(parameters['trade_date'], '%Y-%m-%d').date()
        screen_lf = lf_history.with_columns(pl.lit(None).cast(pl.String).alias('time'))
    else:
        quotation_lf = STK_QUOTATION.output_eager_lazyframe()
        trade_date = quotation_lf.select(pl.col('trade_date').max()).collect().item()
        lf_realtime = (
            quotation_lf
            .with_columns(
                pl.lit(trade_date).cast(pl.Date).alias('trade_date'),
                pl.lit(False).alias('suspension'),
                pl.col('time').cast(pl.String).alias('time'),
            )
            .select([*STK_BUILTIN_FACTORS, 'time'])
        )
        screen_lf = pl.concat([
            (
                lf_history
                    .filter(pl.col('trade_date') < trade_date)
                    .with_columns(pl.lit(None).cast(pl.String).alias('time'))
            ),
            lf_realtime,
        ])
    # 并表：在 concat 后执行，确保历史和实时数据都获得 Join 列
    screen_lf, _builtin_factors = maybe_resolve_joins(screen_lf, 2, parameters, STK_BUILTIN_FACTORS)
    # 选股
    screen_results = stk_screen(
        screen_lf,
        trade_date,
        list(parameters['strategy_configs'].values()),
        parameters.get('null_score_type', NullScoreType.EXCLUDE) == NullScoreType.EXCLUDE,
        parameters.get('show_all_data', False),
        parameters.get('hold_securities'),
        builtin_factors=_builtin_factors,
    )
    answer = dict(
        zip(
            parameters['strategy_configs'].keys(), 
            map(lambda df: df.to_dicts(), screen_results)
        )
    )
    return response('screen', answer)
