from typing import Any

from kit.kit import DataHelper
from strategy import trade_strategy
from task.base_task import base_task


class process_pending_buy_task(base_task):
    """
    处理待买入股票队列
    """

    def __init__(self, strategy: trade_strategy):
        super().__init__(strategy)

    def config(self, context: Any):
        self.today_trade_switch = self.get_config('today_trade_switch')
        self.stock_num = self.get_config('stock_num')
        self.buy_requests_list = self.strategy.state.get_buy_requests()

    def run(self, context: Any):
        # 如果在特殊月份（1月/4月）或没有待买入股票，直接返回
        if self.today_trade_switch or not self.buy_requests_list:
            return

        self.log.info("开始处理待买入股票队列")
        current_data = DataHelper.get_current_data()

        # 更新当前持仓列表
        self.hold_list = [position.security for position in list(context.portfolio.positions.values())]
        position_count = len(self.hold_list)
        target_num = int(self.stock_num)

        # 如果持仓已满，不再买入
        if position_count >= target_num:
            self.log.info("当前持仓已满，清空待买入队列")
            self.strategy.state.clear_buy_requests()
            return

        # 计算可买入的股票数量和分配资金
        buy_count = min(len(self.buy_requests_list), target_num - position_count)
        if buy_count <= 0:
            return

        # 计算每只股票可分配的资金
        try:
            value = context.portfolio.cash / buy_count
        except ZeroDivisionError:
            self.log.error("资金分摊时除零错误")
            return

        # 确保每只股票有足够的买入金额
        if value < 5000:  # 假设最小买入金额为5000元
            self.log.warning(f"每只股票分配资金不足: {value:.2f}元，取消本次买入")
            return

        # 逐个买入股票
        success_count = 0
        for idx, stock in enumerate(self.buy_requests_list[:buy_count]):
            # 跳过已持仓的股票
            if stock in self.hold_list:
                continue

            # 检查是否可交易
            if current_data[stock].paused:
                self.log.warning(f"股票 {stock} 处于停牌状态，跳过买入")
                continue

            if current_data[stock].high_limit <= current_data[stock].last_price:
                self.log.warning(f"股票 {stock} 已涨停，跳过买入")
                continue

            # 执行买入
            if self.strategy.open_position(stock, value):
                self.log.info(f"成功买入股票 {stock}，分配资金 {value:.2f}")
                # self.position_manager.get_not_buy_again().append(stock)
                success_count += 1

                # 如果持仓已满，停止买入
                if len(context.portfolio.positions) >= target_num:
                    break

        # 清理已处理的股票
        # self.position_manager.pending_buys = self.position_manager.pending_buys[buy_count:]

        # self.log.info(f"本次共成功买入 {success_count} 只股票，待买入队列还剩 {len(self.position_manager.pending_buys)} 只")

        # 输出最新持仓状态
        # self.position_manager.log_status(context, [p.security for p in list(context.portfolio.positions.values())])

    def handle(self, context: Any):
        pass

    def end(self, context: Any):
        pass