"""FINANCE Scope 自定义因子（Link Factor）调试 handler"""
from io import BytesIO

import polars as pl

from ldnode.loader.files import STK_FINA_IND
from ldnode.calculator.plans.polars import planning_computed_factors
from .models import FormulaComputeParameters, TaskResponse
from .utils import wrap_handler, response


@wrap_handler("财报公式计算")
def handle_finance_formula_compute(parameters: FormulaComputeParameters) -> TaskResponse[bytes]:
    root_cols = {'code', 'report_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 = STK_FINA_IND.output_lazyframe()
    builtin_factors = list(lf.collect_schema().names())
    df = (
        lf
        .pipe(
            planning_computed_factors,
            builtin_factors=builtin_factors,
            computed_factors=parameters['formula_list'],
            exclude_suspension=False,
            root_cols=root_cols,
        )
        .select(root_cols | customized_factors)
    ).collect()
    pq_buffer = BytesIO()
    df.write_parquet(pq_buffer)
    return response(data_type="formula_compute", data=pq_buffer.getvalue())
