from io import BytesIO

from ldnode.loader.persistence import get_clickhouse_client
from ldnode.calculator.operators import cb_grid_backtest
from .models import CBGridBacktestParameters, CBGridBacktestResponse, TaskResponse
from .utils import wrap_handler, response, error_response


@wrap_handler("转债网格回测")
def handle_cb_grid_backtest(parameters: CBGridBacktestParameters) -> TaskResponse[CBGridBacktestResponse]:
    """ 响应转债网格回测任务 """
    # 获取Clickhouse客户端
    clickhouse_client = get_clickhouse_client()
    # 执行查询
    df = clickhouse_client.query_dataframe(
        """
            select open, toTimeZone(datetime, 'Asia/Shanghai') as datetime, close from cb_data_1m
            where symbol=%(code)s
            and (datetime between %(start_date)s and %(end_date)s)
            order by datetime
        """,
        {
            'code': parameters['strategy_config']['code'], 
            'start_date': parameters['strategy_config']['start_date'], 
            'end_date': parameters['strategy_config']['end_date']
        }
    )
    if df.empty:
        return error_response("所选时间段内没有发生任何数据，请重新选择时间段或标的")
    df = cb_grid_backtest(
        df=df,
        strategy_config=parameters['strategy_config']
    )
    evaluate_dict = {
        'start_time': df.index[0].to_pydatetime().strftime('%Y-%m-%d %H:%M'), 
        'start_pos': df.iat[0, df.columns.get_loc('pos')] if df.iat[0, df.columns.get_loc('pos')] else 0,
        'start_pos_val': df.iat[0, df.columns.get_loc('pos_val')], 
        'start_cash': df.iat[0, df.columns.get_loc('cash')],
        'end_time': df.index[-1].to_pydatetime().strftime('%Y-%m-%d %H:%M'), 
        'end_pos': df.iat[-1, df.columns.get_loc('pos')], 
        'end_pos_val': df.iat[-1, df.columns.get_loc('pos_val')], 
        'end_cash': df.iat[-1, df.columns.get_loc('cash')], 
        'end_asset': df.iat[-1, df.columns.get_loc('asset')],
        'times': f'{(df.index[-1] - df.index[0]).days}天', 
        'cash_mean': df.cash.mean(), 
        'cash_max': df.cash.min(), 
        'grid_return_acc': df.iat[-1, df.columns.get_loc('grid_return_acc')]
    }
    df = df.reset_index()[
        ['datetime','open','close', 'grid_index', 'pos','pos_val','cash','grid_return','grid_return_acc','asset']
    ]
    pq_buffer = BytesIO()
    df.to_parquet(pq_buffer, index=True)
    return response(data_type="grid_backtest_result", data={
        "evaluation": evaluate_dict,
        "history_position": pq_buffer.getvalue(),
    })
