import pymysql
from datetime import datetime
import json
import os
from contextlib import contextmanager
from dotenv import load_dotenv

load_dotenv()


class TradingDatabase:
    """交易数据库管理类 - 用于记录交易历史和预测结果 (MySQL)"""
    
    def __init__(self, db_path=None):
        """初始化 (db_path 参数已废弃，仅为兼容性保留)"""
        self._init_db()

    @contextmanager
    def get_connection(self):
        """获取 MySQL 数据库连接"""
        conn = pymysql.connect(
            host=os.getenv('MYSQL_HOST'),
            port=int(os.getenv('MYSQL_PORT', 3306)),
            user=os.getenv('MYSQL_USER'),
            password=os.getenv('MYSQL_PASSWORD'),
            database=os.getenv('MYSQL_DB'),
            charset='utf8mb4',
            cursorclass=pymysql.cursors.DictCursor
        )
        try:
            yield conn
        finally:
            conn.close()

    def _init_db(self):
        """初始化数据库表结构"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS trades (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    timestamp DATETIME,
                    market_title VARCHAR(255),
                    market_slug VARCHAR(255),
                    model_name VARCHAR(50),
                    prediction VARCHAR(10),
                    probability DOUBLE,
                    action VARCHAR(10),
                    price DOUBLE,
                    amount DOUBLE,
                    size DOUBLE,
                    status VARCHAR(20),
                    error_message TEXT,
                    raw_data JSON,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    order_id VARCHAR(100),
                    INDEX idx_market_slug (market_slug),
                    INDEX idx_status (status)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
            ''')
            conn.commit()
            
            # 检查并添加 order_id 列 (Migration)
            self._ensure_column_exists(conn, 'trades', 'order_id', 'VARCHAR(100)')

    def _ensure_column_exists(self, conn, table_name, column_name, column_type):
        """确保表中存在指定列"""
        try:
            cursor = conn.cursor()
            cursor.execute(f"SHOW COLUMNS FROM {table_name} LIKE '{column_name}'")
            if not cursor.fetchone():
                print(f"正在添加列 {column_name} 到表 {table_name}...")
                cursor.execute(f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_type}")
                conn.commit()
        except Exception as e:
            print(f"检查/添加列失败: {e}")

    def log_trade(self, trade_data: dict):
        """记录交易信息"""
        try:
            with self.get_connection() as conn:
                cursor = conn.cursor()
                
                timestamp = trade_data.get('timestamp', datetime.now().isoformat())
                raw_data_json = json.dumps(trade_data.get('raw_data', {}), ensure_ascii=False)
                
                sql = '''
                    INSERT INTO trades (
                        timestamp, market_title, market_slug, model_name,
                        prediction, probability, action, price, amount, size,
                        status, strategy, error_message, raw_data, order_id
                    ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                '''
                
                cursor.execute(sql, (
                    timestamp,
                    trade_data.get('market_title'),
                    trade_data.get('market_slug'),
                    trade_data.get('model_name'),
                    trade_data.get('prediction'),
                    trade_data.get('probability'),
                    trade_data.get('action'),
                    trade_data.get('price'),
                    trade_data.get('amount'),
                    trade_data.get('size'),
                    trade_data.get('status'),
                    trade_data.get('strategy'),
                    trade_data.get('error_message'),
                    raw_data_json,
                    trade_data.get('order_id')
                ))
                conn.commit()
        except Exception as e:
            print(f"写入数据库失败: {e}")

    def check_existing_trade(self, market_slug: str, status: str = 'success') -> bool:
        """检查指定市场是否已有成功交易记录"""
        try:
            with self.get_connection() as conn:
                cursor = conn.cursor()
                sql = 'SELECT COUNT(*) as cnt FROM trades WHERE market_slug = %s AND status = %s'
                cursor.execute(sql, (market_slug, status))
                result = cursor.fetchone()
                count = result['cnt'] if result else 0
                return count > 0
        except Exception as e:
            print(f"检查已存在交易失败: {e}")
            return False

    def get_existing_trade(self, market_slug: str) -> dict:
        """获取指定市场的最近一次交易记录"""
        try:
            with self.get_connection() as conn:
                cursor = conn.cursor()
                # Get the latest trade for this market
                sql = 'SELECT * FROM trades WHERE market_slug = %s ORDER BY id DESC LIMIT 1'
                cursor.execute(sql, (market_slug,))
                return cursor.fetchone()
        except Exception as e:
            print(f"获取交易记录失败: {e}")
            return None

    def update_trade_status(self, trade_id: int, status: str, error_message: str = None):
        """更新交易状态"""
        try:
            with self.get_connection() as conn:
                cursor = conn.cursor()
                if error_message:
                    sql = 'UPDATE trades SET status = %s, error_message = %s WHERE id = %s'
                    cursor.execute(sql, (status, error_message, trade_id))
                else:
                    sql = 'UPDATE trades SET status = %s WHERE id = %s'
                    cursor.execute(sql, (status, trade_id))
                conn.commit()
        except Exception as e:
            print(f"更新交易状态失败: {e}")