from io import BytesIO

import polars as pl

from ldnode.loader.files import CB_DATA
from ldnode.calculator.schemas import CB_BUILTIN_FACTORS
from ldnode.calculator.plans.polars import maybe_resolve_joins, planning_computed_factors
from .models import FormulaComputeParameters, TaskResponse
from .utils import wrap_handler, response


@wrap_handler("可转债公式计算")
def handle_cb_formula_compute(parameters: FormulaComputeParameters) -> TaskResponse[bytes]:
    """ 响应可转债公式计算任务 """
    root_cols = {'code', 'name', 'trade_date'}
    # 获取历史数据 
    customized_factors = set()
    for fs in parameters['formula_list']:
        customized_factors.update([f['factor'] for f in fs if not f['is_intermediate']])
    # 获取历史数据
    lf = (
        CB_DATA.output_lazyframe()
        .filter(~pl.col('suspension'))
        # 补丁
        .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')
        )
    )
    lf, _builtin_factors = maybe_resolve_joins(lf, 1, parameters, CB_BUILTIN_FACTORS)
    df = (
        lf
        # .filter(col('_fresh_index') <= 2 * max_window) # 保留最大_fresh_index窗口
        # 计算自定义因子并追加其根列到root_cols
        .pipe(
            planning_computed_factors,
            builtin_factors=_builtin_factors,
            computed_factors=parameters['formula_list'],
            exclude_suspension=False, 
            root_cols=root_cols
        )
        # .filter(col('_fresh_index') <= max_window)
        .filter(pl.col('trade_date') == lf.select(pl.col('trade_date')).max().collect().item())
        .select(root_cols | customized_factors)
    ).collect()
    pq_buffer = BytesIO()
    df.write_parquet(pq_buffer)
    return response(data_type="formula_compute", data=pq_buffer.getvalue())
