# filename: tqsdk_context.py

from datetime import datetime
from typing import Optional, Any, Dict, List, Literal, Deque, TYPE_CHECKING
from collections import deque

# 导入你提供的 core_data 中的类型
from src.core_data import Bar, Order, Trade, PortfolioSnapshot # 确保此路径正确，如果core_data不在同级目录，需要调整

# 导入 Tqsdk 的核心类型
import tqsdk
from tqsdk import TqApi, TqAccount, tafunc

import pandas as pd

# 使用 TYPE_CHECKING 避免循环导入，只在类型检查时导入 TqsdkEngine
if TYPE_CHECKING:
    from src.tqsdk_engine import TqsdkEngine # 假设 TqsdkEngine 在 tqsdk_engine.py 中


class TqsdkContext:
    """
    Tqsdk 回测上下文，适配原有 BacktestContext 接口。
    策略通过此上下文与 Tqsdk 进行交互。
    """
    def __init__(self, api: TqApi):
        """
        初始化 Tqsdk 回测上下文。

        Args:
            api (TqApi): Tqsdk 的 TqApi 实例。
        """
        self._api = api
        self._current_bar: Optional[Bar] = None
        self._engine: Optional['TqsdkEngine'] = None # 添加对引擎的引用，用于访问其状态或触发事件
        
        # 用于缓存 Tqsdk 的 K 线序列，避免每次都 get_kline_serial
        self._kline_serial: Dict[str, object] = {}

        # 订单/取消请求队列，TqsdkEngine 会在异步循环中处理它们
        self.order_queue: Deque[Order] = deque()
        self.cancel_queue: Deque[str] = deque() # 存储 order_id

        print("TqsdkContext: 初始化完成。")

    def set_current_bar(self, bar: Bar):
        """
        设置当前正在处理的 K 线数据。
        由 TqsdkEngine 调用。
        """
        self._current_bar = bar

    def get_current_bar(self) -> Optional[Bar]:
        """
        获取当前正在处理的 K 线数据。
        策略可以通过此方法获取最新 K 线。
        """
        return self._current_bar

    def get_kline_data(self, symbol: str, duration_seconds: int, data_length: int = 10):
        """
        获取指定合约的 K 线数据。
        返回 Tqsdk 的 DataFrame 格式 K 线序列（TqKLine 对象），可以直接用于计算指标。
        如果需要转换为你自己的 Bar 对象列表，则需要在此方法内部进行转换。
        """
        if symbol not in self._kline_serial:
            # 这里的 get_kline_serial 并不是实时获取，而是在 TqApi 启动时就已经加载
            # 所以在 Context 中直接调用是安全的，TqApi 会返回已加载的数据引用
            self._kline_serial[symbol] = self._api.get_kline_serial(symbol, duration_seconds, data_length=data_length)
        return self._kline_serial[symbol]

    def get_current_time(self) -> datetime:
        """
        获取当前模拟时间（Tqsdk 的数据时间）。
        """
        # Tqsdk 的 get_tick_timestamp() 返回微秒时间戳
        return self.get_current_bar().datetime

    def get_current_positions(self) -> Dict[str, int]:
        """
        获取当前所有持仓。返回 {symbol: quantity} 的字典，quantity 为净持仓量（多头-空头）。
        """
        tq_positions: Dict[str] = self._api.get_position()
        converted_positions: Dict[str, int] = {}
        for symbol, pos in tq_positions.items():
            net_pos = pos.pos_long - pos.pos_short
            if net_pos != 0:
                converted_positions[symbol] = net_pos
        return converted_positions
    
    def get_pending_orders(self) -> Dict[str, Order]:
        """
        获取当前所有待处理（未成交）订单。
        返回 {order_id: Order} 的字典。
        """
        tq_orders: Dict[str] = self._api.get_order()
        pending_orders: Dict[str, Order] = {}
        for order_id, tq_order in tq_orders.items():
            if tq_order.status == "ALIVE": # 正在进行中的订单
                # 将 TqOrder 转换为你自己的 core_data.Order 类型
                # 注意：core_data.Order 的 direction 有 "CLOSE_LONG", "CLOSE_SHORT" 等，需要映射
                # Tqsdk 的 direction 只有 "BUY", "SELL"
                # Tqsdk 的 offset 决定了是开仓还是平仓
                core_direction: Literal["BUY", "SELL", "CLOSE_LONG", "CLOSE_SHORT"]
                if tq_order.offset == "OPEN":
                    core_direction = tq_order.direction # 开仓时方向直接对应买卖
                elif tq_order.offset in ["CLOSE", "CLOSETODAY", "CLOSEYESTERDAY"]:
                    # 平仓时，买入平空，卖出平多
                    core_direction = "CLOSE_SHORT" if tq_order.direction == "BUY" else "CLOSE_LONG"
                else: # 默认为 BUY/SELL
                    core_direction = tq_order.direction


                converted_order = Order(
                    id=tq_order.order_id, # 将 Tqsdk 的 order_id 赋值给你的 Order 类的 id
                    symbol=tq_order.exchange_id + "." + tq_order.instrument_id, # 例如 "SHFE.rb2401"
                    direction=core_direction,
                    volume=tq_order.volume_orign,
                    price_type="LIMIT" if tq_order.limit_price is not None else "MARKET", # Tqsdk 市价单类型为 "ANY"
                    limit_price=tq_order.limit_price,
                    offset=tq_order.offset, # Tqsdk 原生 offset
                    # order_id=tq_order.order_id, # 存储 Tqsdk 的 order_id
                    submitted_time=pd.to_datetime(tq_order.insert_date_time, unit="ns", utc=True),
                    # status=tq_order.status # 保持 Tqsdk 的状态字符串
                )
                pending_orders[order_id] = converted_order
        return pending_orders

    def get_account_cash(self) -> float:
        """
        获取当前可用现金。
        """
        account: TqAccount = self._api.get_account()
        return account.available_cash if account else 0.0

    def get_average_position_price(self, symbol: str) -> Optional[float]:
        """
        获取指定合约的平均持仓成本。
        注意: Tqsdk 的 TqPosition 对象中包含了 open_price_long 和 open_price_short。
        这里需要根据多头或空头持仓返回对应的平均成本。
        """
        position = self._api.get_position(symbol)
        # if position:
        #     return avg_cost
        if position:
            if position.pos_long > 0:
                return position.open_price_long
            elif position.pos_short > 0:
                return position.open_price_short
        return None

    def send_order(self, order: Order) -> Optional[Order]:
        """
        策略通过此方法发送订单。
        将订单放入队列，等待 TqsdkEngine 在其异步循环中处理。
        """
        # 为订单分配一个临时ID，便于在队列中追踪，实际ID由Tqsdk返回后更新
        if not order.id: # 使用 order.id 属性
            order.id = f"LOCAL_{id(order)}_{datetime.now().strftime('%f')}"
            order.order_id = order.id # 保持 Tqsdk 风格的 order_id 也一致
        self.order_queue.append(order)
        print(f"Context: 订单已加入队列: {order}")
        return order # 返回传入的订单，待引擎更新其状态和ID

    def cancel_order(self, order_id: str) -> bool:
        """
        策略通过此方法取消指定ID的订单。
        将取消请求放入队列，等待 TqsdkEngine 在其异步循环中处理。
        """
        # 检查订单是否处于待处理状态
        if order_id in self.get_pending_orders(): 
            self.cancel_queue.append(order_id)
            print(f"Context: 取消订单请求已加入队列: {order_id}")
            return True
        print(f"Context: 订单 {order_id} 不在待处理队列中，无法取消。")
        return False

    def set_engine(self, engine: 'TqsdkEngine'): # 使用 TYPE_CHECKING 中的 TqsdkEngine 类型提示
        """
        设置对 TqsdkEngine 实例的引用。
        由 TqsdkEngine 在初始化时调用，用于允许 Context 访问 Engine 的状态。
        """
        self._engine = engine
        print("TqsdkContext: 已设置引擎引用。")

    @property
    def is_rollover_bar(self) -> bool:
        """
        属性：判断当前 K 线是否为换月 K 线（即新合约的第一根 K 线）。
        用于在换月时禁止策略开仓。
        Tqsdk 的回测模式下，通常通过主力连续合约或多合约同时回测来处理换月。
        此处为适配原有接口的简化实现。如果你需要 Tqsdk 的换月逻辑，
        可能需要在 TqsdkEngine 中实现更复杂的判断，并通过 Context 暴露此状态。
        对于 Tqsdk 的主力连续合约，通常不需要策略层面关心具体的换月 K 线。
        """
        # 如果引擎设置了 is_rollover_bar 属性，则使用引擎的判断
        if self._engine and hasattr(self._engine, 'is_rollover_bar'):
            return self._engine.is_rollover_bar
        return False # 默认返回 False
    
    def get_bar_history(self):
        return self._engine.get_bar_history()