from typing import Any
from datetime import datetime

import polars as pl

from ldnode.loader.files import ETF_DATA, ETF_QUOTATION
from ldnode.calculator.schemas import ETF_BUILTIN_FACTORS
from ldnode.calculator.operators import etf_screen
from .utils import wrap_handler, response
from .models import ETFScreenParameters, TaskResponse, NullScoreType


@wrap_handler("选基任务")
def handle_etf_screen(parameters: ETFScreenParameters) -> TaskResponse[dict[str, list[dict[str, Any]]]]:
    """ 响应选基任务 """
    # ETF 数据无 _fresh_index 列，不支持预切片优化
    lf_history = ETF_DATA.output_lazyframe()
    lf_history = lf_history.filter(~pl.col('suspension')).select(ETF_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:
        # 实时选基：从实时行情的 trade_date 字段获取真实交易日，
        # 交易日盘中 = 今天，非交易日 = 上一交易日，避免与历史截面数据重叠
        quotation_lf = ETF_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([*ETF_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,
        ])
    # 选债
    screen_results = etf_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'),
    )
    answer = dict(
        zip(
            parameters['strategy_configs'].keys(), 
            map(lambda df: df.to_dicts(), screen_results)
        )
    )
    return response('screen', answer)
    