import polars as pl

from .models import (
    TaskResponse,
    GenBasketParameters,
    GenBasketResultItem,
)
from .utils import (
    wrap_handler,
    response,
)


@wrap_handler("生成篮子")
def handle_basket_gen_task(parameters: GenBasketParameters) -> TaskResponse[list[GenBasketResultItem]]:
    """ 响应篮子生成任务 """
    # 生成计算计划
    plan = (
        pl.LazyFrame(parameters['basket_data'])
            .select(
                pl.col('code'),
                (
                    parameters['capital'] 
                    * (pl.col('weight') / pl.col('weight').sum()) 
                    // (pl.col('price') * parameters['board_lot'])
                    * parameters['board_lot']
                ).cast(int).alias('target_num'),
            )
    )
    baskets = plan.collect().to_dicts()
    return response("basket", baskets)
