from io import BytesIO
from typing import Any, Generator

import polars as pl

from ldnode import logger
from ldnode.calculator.operators import cb_event_backtest, cb_vector_backtest, evaluate
from ldnode.calculator.plans.polars import (
    maybe_resolve_joins,
    planning_earnings_curve,
    planning_merge_benchmark,
)
from ldnode.calculator.schemas import CB_BUILTIN_FACTORS
from ldnode.loader.files import CB_DATA

from .models import BacktestType, CBBacktestParameters, TaskResponse
from .utils import Timer, get_index_data_lf, response, wrap_handler


@wrap_handler("转债回测", is_generator=True)
def handle_cb_backtest(
    parameters: CBBacktestParameters,
) -> Generator[TaskResponse[Any]]:
    """响应转债回测任务"""
    # 获取历史数据
    lf = (
        CB_DATA.output_lazyframe()
        # 补丁
        .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)
    res_df = None
    history_position = {}
    # 回测
    with Timer(
        f"转债回测结束. 任务ID: {parameters['task_id']} ", logger=logger.event_backtest
    ):
        match parameters["strategy_config"]["backtest_type"]:
            case BacktestType.VECTOR.value:
                res_df = cb_vector_backtest(
                    lf=lf,
                    parameters=parameters["strategy_config"],
                    builtin_factors=_builtin_factors,
                )
                history_position = pl.DataFrame()
            case BacktestType.EVENT.value:
                res_df, history_position = cb_event_backtest(
                    lf=lf,
                    parameters=parameters["strategy_config"],
                    builtin_factors=_builtin_factors,
                    record_history=True,
                )
            case _:
                raise ValueError(
                    f"未知的回测类型: {parameters['strategy_config']['backtest_type']}"
                )
    # 合并基准benchmark
    with Timer(
        f"计算收益曲线结束. 任务ID: {parameters['task_id']} ",
        logger=logger.event_backtest,
    ):
        benchmark_lf = get_index_data_lf(parameters["strategy_config"]["benchmark"])
        merged_lf = planning_merge_benchmark(res_df.lazy(), benchmark_lf)
        res_df = planning_earnings_curve(
            merged_lf, parameters["strategy_config"]["capital"]
        ).collect()
        result_buffer = BytesIO()
        res_df.write_parquet(result_buffer)
        # 解析 COS 存储路径（优先使用调用方预生成的路径，否则使用 task_id 生成默认路径）
        _sk = parameters.get("storage_keys") or {}
        result_key = _sk.get("result_key") or f"backtest/{parameters['task_id']}/timeline.parquet"
        evalution_key = _sk.get("evalution_key") or f"backtest-evalution/{parameters['task_id']}.json"
        position_key = _sk.get("position_key") or f"backtest/{parameters['task_id']}/history_position.parquet"
        # 返回收益序列
        yield response(
            data_type="result",
            data=result_buffer.getvalue(),
            storage_key=result_key,
        )
    # 生成评价指标
    with Timer(
        f"生成评价指标结束. 任务ID: {parameters['task_id']} ",
        logger=logger.event_backtest,
    ):
        evaluate_dict = {
            "strategy": evaluate(res_df),
            "benchmark": evaluate(
                res_df,
                columns_dict={
                    "equity": "equity_benchmark",
                    "time_return": "time_return_benchmark",
                    "drawdown": "drawdown_benchmark",
                },
            ),
            "relative_excess": evaluate(
                res_df,
                columns_dict={
                    "equity": "equity_relative_excess",
                    "time_return": "time_return_relative_excess",
                    "drawdown": "drawdown_relative_excess",
                },
            ),
            "absolute_excess": evaluate(
                res_df,
                columns_dict={
                    "equity": "equity_absolute_excess",
                    "time_return": "time_return_absolute_excess",
                    "drawdown": "drawdown_absolute_excess",
                },
            ),
        }
        yield response(
            data_type="evalution",
            data=evaluate_dict,
            storage_key=evalution_key,
        )
    # 上传每日持仓
    history_postion_buffer = BytesIO()
    history_position.write_parquet(history_postion_buffer, row_group_size=5000)
    yield response(
        data_type="history_position",
        data=history_postion_buffer.getvalue(),
        storage_key=position_key,
    )
