from typing import Any
from datetime import datetime

import polars as pl

from ldnode.loader.files import CB_DATA, CB_QUOTATION, STK_DATA, STK_QUOTATION
from ldnode.calculator.schemas import CB_BUILTIN_FACTORS
from ldnode.calculator.operators import cb_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 .models import (
    TaskResponse,
    CBScreenParameters,
    NullScoreType,
)
from .utils import wrap_handler, response


@wrap_handler("选债任务")
def handle_cb_screen(parameters: CBScreenParameters) -> TaskResponse[dict[str, list[dict[str, Any]]]]:
    """ 响应选债任务 """
    # 历史数据预切片：按自定义因子最大时间窗口过滤 _fresh_index，减少滚动计算量
    max_window = max(
        (extract_max_window(cfg.get("computed_factors"))
         for cfg in parameters['strategy_configs'].values()),
        default=0,
    )
    lf_history = CB_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(CB_BUILTIN_FACTORS)
        # 补丁
        .with_columns(
            pl.col('code').cast(pl.String).alias('code'),
            pl.col('limit').cast(pl.Int8).alias('limit'),
            pl.col('rating').cast(pl.String).alias('rating'),
            pl.col('yy_rating').cast(pl.String).alias('yy_rating'),
            pl.col('redeem_remain_days').cast(pl.Int64).alias('redeem_remain_days'),
        )
    )
    # 处理选债截面
    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 = CB_QUOTATION.output_eager_lazyframe()
        trade_date = quotation_lf.select(pl.col('date').cast(pl.Date).max()).collect().item()
        lf_realtime = (
            quotation_lf
            .with_columns(
                pl.col('code').cast(pl.String).alias('code'),
                pl.lit(trade_date).cast(pl.Date).alias('trade_date'),
                pl.col('vol_stk').cast(pl.Float64).alias('vol_stk'),
                pl.col('rating').cast(pl.String).alias('rating'),
                pl.col('yy_rating').cast(pl.String).alias('yy_rating'),
                pl.col('redeem_remain_days').cast(pl.Int64).alias('redeem_remain_days'),
                pl.col('limit').cast(pl.Int8).alias('limit'),
                pl.lit(False).alias('suspension'),
                pl.col('time').cast(pl.String).alias('time'),
                pl.col('is_call').cast(pl.String).alias('is_call')
            )
            .select([*CB_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 不到
    aux_data_overrides = None
    if not parameters.get('trade_date'):
        # CONCAT(STK_DATA, STK_QUOTATION) 作为正股附表，覆盖默认的 STK_DATA
        stk_quotation_lf = STK_QUOTATION.output_eager_lazyframe().with_columns(
            pl.col('trade_date').cast(pl.Date),
        )
        aux_data_overrides = {
            2: pl.concat([STK_DATA.output_lazyframe(), stk_quotation_lf], how="diagonal"),
        }
    screen_lf, _builtin_factors = maybe_resolve_joins(
        screen_lf, 1, parameters, CB_BUILTIN_FACTORS,
        aux_data_overrides=aux_data_overrides,
    )
    # 选债
    screen_results = cb_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)
