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,
                    INDEX idx_market_slug (market_slug),
                    INDEX idx_status (status)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
            ''')
            conn.commit()

    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
                    ) VALUES (%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
                ))
                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