from typing import Any
from datetime import datetime

import polars as pl

from ldnode import logger
from ldnode.loader.files import CB_DATA, CB_QUOTATION, STK_DATA, STK_QUOTATION
from ldnode.calculator.schemas import CB_BUILTIN_FACTORS, CB_SCREEN_RESULT_FIXED_COLUMNS
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
from .risk_control import (
    cb_risk_control_expr,
    count_risk_control_rows,
    realtime_risk_control_fixed_columns,
)


def concat_history_before_trade_date_with_realtime(
    history_lf: pl.LazyFrame,
    realtime_lf: pl.LazyFrame,
    trade_date,
    module: str,
) -> pl.LazyFrame:
    logger.screen.info(
        f"[{module}] 盘中附表拼接开始 history_filter=trade_date<trade_date realtime_trade_date={trade_date}"
    )
    return pl.concat(
        [
            history_lf.filter(pl.col("trade_date") < trade_date),
            realtime_lf,
        ],
        how="diagonal",
    )


@wrap_handler("选债任务")
def handle_cb_screen(parameters: CBScreenParameters) -> TaskResponse[dict[str, list[dict[str, Any]]]]:
    """ 响应选债任务 """
    is_realtime_screen = not parameters.get('trade_date')
    # 历史数据预切片：按自定义因子最大时间窗口过滤 _fresh_index，减少滚动计算量
    max_window = max(
        (extract_max_window(cfg.get("computed_factors"))
         for cfg in parameters['strategy_configs'].values()),
        default=0,
    )
    logger.screen.info(
        "[CBScreen] 历史预切片决策 max_window=%s buffer=30 limit=%s strategy_count=%s realtime=%s",
        max_window,
        max_window + 30 if max_window > 0 else None,
        len(parameters['strategy_configs']),
        is_realtime_screen,
    )
    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('list_days').cast(pl.Int64).alias('list_days'),
            pl.col('redeem_real_days').cast(pl.Int64).alias('redeem_real_days'),
            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()
        quotation_schema = quotation_lf.collect_schema()
        has_realtime_quote_name = 'name' in quotation_schema.names()
        trade_date = quotation_lf.select(pl.col('date').cast(pl.Date).max()).collect().item()
        logger.screen.info(
            "[RiskControl] 实时风控字段生成开始 method=cb_screen has_delist_date=%s has_realtime_quote_name=%s trade_date=%s",
            False,
            has_realtime_quote_name,
            trade_date,
        )
        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('list_days').cast(pl.Int64).alias('list_days'),
                pl.col('redeem_real_days').cast(pl.Int64).alias('redeem_real_days'),
                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,
        ]).with_columns(cb_risk_control_expr())
    # 并表：盘中模式下附表也需要 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: concat_history_before_trade_date_with_realtime(
                STK_DATA.output_lazyframe(),
                stk_quotation_lf,
                trade_date,
                "CBScreen",
            ),
        }
    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,
        fixed_columns=realtime_risk_control_fixed_columns(CB_SCREEN_RESULT_FIXED_COLUMNS) if is_realtime_screen else None,
    )
    if is_realtime_screen:
        total_rows = sum(result.height for result in screen_results)
        risk_rows = sum(count_risk_control_rows(result) for result in screen_results)
        logger.screen.info(
            "[RiskControl] 实时风控字段生成完成 method=cb_screen result_groups=%s rows=%s risk_rows=%s",
            len(screen_results),
            total_rows,
            risk_rows,
        )
    answer = dict(
        zip(
            parameters['strategy_configs'].keys(), 
            map(lambda df: df.to_dicts(), screen_results)
        )
    )
    return response('screen', answer)
