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()
        self.hold_list = [position.security for position in list(context.portfolio.positions.values())]
        self.buy_symbol_money_up_limit = self.get_config('buy_symbol_money_up_limit')
        self.buy_symbol_money_total = self.get_config('buy_symbol_money_total')

    def run(self, context: Any):
        # 如果在特殊月份（1月/4月）或没有待买入股票，直接返回
        if self.today_trade_switch or not self.buy_requests_list:
            self.log.info("今日非交易日或没有待买入股票，跳过处理")
            return

        self.log.info("开始处理待买入股票队列")
        current_data = DataHelper.get_current_data()
        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:
            self.log.info("没有可买入的股票，或持仓已满")
            return

        # 计算每只股票可分配的资金
        if context.portfolio.cash < self.buy_symbol_money_total:
            self.log.warning(f"可用资金 {context.portfolio.cash:.2f} 元不足，无法满足总资金要求 {self.buy_symbol_money_total} 元")
            return

        try:
            value = context.portfolio.cash / buy_count
        except ZeroDivisionError:
            self.log.error("资金分摊时除零错误")
            return

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

        # 逐个买入股票
        for idx, stock in enumerate(self.buy_requests_list[:buy_count]):
            # 跳过已持仓的股票
            if stock in self.hold_list:
                self.log.warning(f"股票 {stock} 已在持仓中，跳过买入")
                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}")
                # 删除购买列表

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

    def handle(self, context: Any):
        pass

    def end(self, context: Any):
        pass