# coding:utf-8
import datetime
import sys
from xtquant.xttrader import XtQuantTraderCallback
from strategy import update_position_in_memory, remove_pending_order
from database_manager import DatabaseManager

class MyXtQuantTraderCallback(XtQuantTraderCallback):
    def __init__(self, logger):
        self.logger = logger
        self.db_manager = DatabaseManager(logger)

    def on_disconnected(self):
        """
        连接断开
        :return:
        """
        self.logger.warning(f"{datetime.datetime.now()} 连接断开回调")

    def on_stock_order(self, order):
        """
        委托回报推送
        :param order: XtOrder对象
        :return:
        """
        self.logger.info(f"{datetime.datetime.now()} 委托回调 {order.order_remark}")
        
        # 更新在途订单状态和数据库
        try:
            # 从订单备注中提取股票代码
            order_remark = order.order_remark
            if '_' in order_remark:
                stock_code = order_remark.split('_')[1]  # 修正：取第二段作为股票代码
                
                # 检查Redis中是否存在该在途订单
                from strategy import initialize_redis_manager
                r = initialize_redis_manager(self.logger)
                order_info = r.get_pending(stock_code)
                
                if order_info:
                    # 确保 order_id 是字符串类型
                    qmt_order_id_str = str(order.order_id)
                    our_order_id = order_info.get('order_id')
                    
                    # 先更新QMT订单ID映射
                    if our_order_id:
                        self.db_manager.update_qmt_order_id(our_order_id, qmt_order_id_str)
                    
                    # 更新数据库订单状态
                    self.db_manager.update_order_status(qmt_order_id_str, 'submitted')
                    
                    # 记录交易日志
                    log_data = {
                        'order_id': our_order_id or qmt_order_id_str,
                        'stock_code': stock_code,
                        'log_type': 'order_submitted',
                        'log_level': 'INFO',
                        'message': f'订单已提交: {order.order_remark}',
                        'create_time': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                    }
                    self.db_manager.insert_trading_log(log_data)
                    
                    self.logger.info(f"订单状态更新为已提交: {stock_code}")
                else:
                    self.logger.warning(f"未找到在途订单: {stock_code}")
        except Exception as e:
            self.logger.error(f"更新订单状态失败: {str(e)}")

    def on_stock_trade(self, trade):
        """
        成交变动推送
        :param trade: XtTrade对象
        :return:
        """
        self.logger.info(f"{datetime.datetime.now()} 成交回调 {trade.order_remark} 成交价格 {trade.traded_price} 成交数量 {trade.traded_volume}")
        
        # 更新内存中的持仓状态和数据库
        try:
            # 根据订单备注判断是买入还是卖出
            is_buy = True  # 默认买入
            is_clearance = False  # 是否为清仓订单
            if trade.order_remark.startswith('CLEARANCE_'):
                is_buy = False
                is_clearance = True
            elif trade.order_remark.startswith('SELL_'):
                is_buy = False
            elif trade.order_remark.startswith('BUY_'):
                is_buy = True
            
            # 记录交易方向
            trade_direction = "买入" if is_buy else "卖出"
            if is_clearance:
                trade_direction = "清仓"
            self.logger.info(f"识别交易方向: {trade_direction}")
            
            # 更新内存和数据库持仓状态（确保价格类型为float）
            traded_price = float(trade.traded_price)
            update_position_in_memory(trade.stock_code, trade.traded_volume, is_buy, traded_price, self.logger)
            
            # 确保 order_id 是字符串类型
            order_id_str = str(trade.order_id)
            
            # 1. 先查询当前订单的已成交数量和委托数量
            current_filled = self._get_current_filled_quantity(order_id_str)
            order_quantity = self._get_order_quantity(order_id_str)
            
            # 2. 计算新的累计成交数量
            new_total_filled = current_filled + trade.traded_volume
            
            # 3. 判断订单是否完全成交
            is_order_completed = new_total_filled >= order_quantity
            
            # 4. 更新数据库订单状态和成交数量
            if is_order_completed:
                # 完全成交
                self.db_manager.update_order_status(
                    order_id_str, 
                    'completed', 
                    new_total_filled, 
                    datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                )
                self.logger.info(f"订单完全成交: {order_id_str} 累计成交 {new_total_filled}/{order_quantity}")
                
                # 从在途订单中移除
                remove_pending_order(trade.stock_code, self.logger)
                
                # 如果是清仓订单且完全成交，更新清仓请求状态
                if is_clearance:
                    completed_time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                    self.db_manager.update_clearance_status(
                        trade.stock_code,
                        'completed',
                        completed_time
                    )
                    # 从Redis中移除清仓计划
                    from strategy import initialize_redis_manager
                    r = initialize_redis_manager(self.logger)
                    r.del_clearance_plan(trade.stock_code)
                    self.logger.info(f"清仓订单完全成交，已更新清仓请求状态: {trade.stock_code}")
            else:
                # 部分成交
                self.db_manager.update_order_status(
                    order_id_str, 
                    'filled', 
                    new_total_filled, 
                    datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                )
                self.logger.info(f"订单部分成交: {order_id_str} 累计成交 {new_total_filled}/{order_quantity}")
            
            # 5. 记录详细交易日志（包含成交信息）
            trade_detail = {
                'trade_id': getattr(trade, 'trade_id', ''),
                'traded_price': traded_price,
                'traded_volume': int(trade.traded_volume),
                'traded_amount': float(traded_price * trade.traded_volume),
                'trade_time': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'total_filled': new_total_filled,
                'order_quantity': order_quantity,
                'is_completed': is_order_completed,
                'trade_direction': trade_direction
            }
            
            log_data = {
                'order_id': order_id_str,
                'stock_code': trade.stock_code,
                'log_type': 'trade_filled',
                'log_level': 'INFO',
                'message': f'{trade_direction}成交: {trade.stock_code} {trade.traded_volume}股 @ {trade.traded_price}元 (累计: {new_total_filled}/{order_quantity})',
                'extra_data': trade_detail,
                'create_time': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            }
            self.db_manager.insert_trading_log(log_data)
            
        except Exception as e:
            self.logger.error(f"更新持仓状态失败: {str(e)}")
    
    def _get_current_filled_quantity(self, order_id):
        """获取订单当前已成交数量"""
        try:
            sql = "SELECT COALESCE(filled_quantity, 0) FROM trading_order WHERE order_id = %s OR qmt_order_id = %s"
            result = self.db_manager.execute_query(sql, (order_id, order_id))
            return int(result[0][0]) if result else 0
        except Exception as e:
            self.logger.warning(f"获取已成交数量失败: {str(e)}")
            return 0
    
    def _get_order_quantity(self, order_id):
        """获取订单委托数量"""
        try:
            sql = "SELECT order_quantity FROM trading_order WHERE order_id = %s OR qmt_order_id = %s"
            result = self.db_manager.execute_query(sql, (order_id, order_id))
            return int(result[0][0]) if result else 0
        except Exception as e:
            self.logger.warning(f"获取委托数量失败: {str(e)}")
            return 0

    def on_order_error(self, order_error):
        """
        委托失败推送
        :param order_error:XtOrderError 对象
        :return:
        """
        self.logger.error(f"委托报错回调 {order_error.order_remark} {order_error.error_msg}")
        
        # 从在途订单中移除失败的订单，并更新数据库
        try:
            order_remark = order_error.order_remark
            if '_' in order_remark:
                stock_code = order_remark.split('_')[1]  # 修正：取第二段作为股票代码
                
                # 确保 order_id 是字符串类型
                order_id_str = str(order_error.order_id)
                
                # 更新数据库订单状态
                self.db_manager.update_order_status(order_id_str, 'failed')
                
                # 记录错误日志
                log_data = {
                    'order_id': order_id_str,
                    'stock_code': stock_code,
                    'log_type': 'order_failed',
                    'log_level': 'ERROR',
                    'message': f'订单失败: {order_error.error_msg}',
                    'create_time': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                }
                self.db_manager.insert_trading_log(log_data)
                
                remove_pending_order(stock_code, self.logger)
        except Exception as e:
            self.logger.error(f"处理失败订单失败: {str(e)}")

    def on_cancel_error(self, cancel_error):
        """
        撤单失败推送
        :param cancel_error: XtCancelError 对象
        :return:
        """
        self.logger.error(f"{datetime.datetime.now()} {sys._getframe().f_code.co_name}")

    def on_order_stock_async_response(self, response):
        """
        异步下单回报推送
        :param response: XtOrderResponse 对象
        :return:
        """
        self.logger.info(f"异步委托回调 {response.order_remark}")

    def on_cancel_order_stock_async_response(self, response):
        """
        :param response: XtCancelOrderResponse 对象
        :return:
        """
        self.logger.info(f"{datetime.datetime.now()} {sys._getframe().f_code.co_name}")

    def on_account_status(self, status):
        """
        :param response: XtAccountStatus 对象
        :return:
        """
        self.logger.info(f"{datetime.datetime.now()} {sys._getframe().f_code.co_name}") 