import os

file_path = 'trader/auto_trade_polymarket.py'

with open(file_path, 'r', encoding='utf-8') as f:
    lines = f.readlines()

new_lines = []
inserted = False

# Code to insert
insert_code = """
        # === 新增: 保存初始预测和状态 ===
        try:
            # 确定最终预测 (用于日志)
            final_pred = 0
            if 'ensemble' in all_preds and all_preds['ensemble']:
                final_pred = all_preds['ensemble']['prediction']
            elif 'dlinear' in all_preds:
                final_pred = all_preds['dlinear']['prediction']
            
            et_tz = ZoneInfo('America/New_York')
            current_time_utc = pd.Timestamp(current_time)
            current_time_et = current_time_utc.tz_localize('UTC').tz_convert(et_tz)
            
            # 1. 保存预测到 model_predictions
            common_data = {
                "timestamp": current_time_utc.isoformat(),
                "timestamp_et": current_time_et.isoformat(),
                "date": current_time_et.strftime('%Y-%m-%d'),
                "hour": current_time_et.hour,
                "actual_result": None,
                "is_correct": None,
                "open_price": float(df['open'].iloc[-1]),
                "close_price": None,
            }
            
            def save_single_pred(name, pred_res):
                data = common_data.copy()
                data['model_name'] = name
                data['prediction'] = pred_res['prediction']
                data['probability'] = pred_res.get('probability', 0.5)
                PredictionModel.save_prediction(data)

            if 'dlinear' in all_preds: save_single_pred("DLinear", all_preds['dlinear'])
            if 'nbeats' in all_preds: save_single_pred("N-BEATS", all_preds['nbeats'])
            if 'ttm' in all_preds: save_single_pred("TTM-R2", all_preds['ttm'])
            if 'ensemble' in all_preds and all_preds['ensemble']: save_single_pred("Ensemble", all_preds['ensemble'])
            
            logger.info(f"💾 [Init] 初始预测已保存到数据库")

            # 2. 记录当前监控状态到 trades (作为 'in_progress')
            status = 'in_progress'
            market_title = self.current_market['title'] if self.current_market else "Waiting for Market..."
            market_slug = self.current_market['slug'] if self.current_market else None
            
            log_data = {
                'timestamp': datetime.now().isoformat(),
                'market_title': market_title,
                'market_slug': market_slug,
                'model_name': 'Ensemble',
                'prediction': 'UP' if final_pred == 1 else 'DOWN',
                'probability': 0.0,
                'action': 'MONITOR',
                'price': 0,
                'amount': 0,
                'size': 0,
                'status': status,
                'error_message': 'System Started - Monitoring',
                'raw_data': {}
            }
            self.db.log_trade(log_data)
            logger.info(f"💾 [Init] 初始监控状态已记录")

        except Exception as e:
            logger.error(f"保存初始状态失败: {e}")
            import traceback
            traceback.print_exc()
            
"""

for line in lines:
    if "def wait_for_next_hour(self)" in line and not inserted:
        new_lines.append(insert_code)
        inserted = True
    new_lines.append(line)

with open(file_path, 'w', encoding='utf-8') as f:
    f.writelines(new_lines)

print("Successfully patched trader/auto_trade_polymarket.py")