import sqlite3
from datetime import datetime
import json
import os
from contextlib import contextmanager

class TradingDatabase:
    """交易数据库管理类 - 用于记录交易历史和预测结果"""
    
    def __init__(self, db_path="trading_history.db"):
        self.db_path = db_path
        self._init_db()

    @contextmanager
    def get_connection(self):
        """数据库连接上下文管理器"""
        conn = sqlite3.connect(self.db_path)
        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 INTEGER PRIMARY KEY AUTOINCREMENT,
                    timestamp TEXT,
                    market_title TEXT,
                    market_slug TEXT,
                    model_name TEXT,
                    prediction TEXT,
                    probability REAL,
                    action TEXT,
                    price REAL,
                    amount REAL,
                    size REAL,
                    status TEXT,
                    error_message TEXT,
                    raw_data TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            conn.commit()

    def log_trade(self, trade_data: dict):
        """
        记录交易信息
        
        Args:
            trade_data: 包含交易信息的字典，支持字段:
                timestamp, market_title, market_slug, model_name,
                prediction, probability, action, price, amount, size,
                status, error_message, raw_data
        """
        try:
            with self.get_connection() as conn:
                cursor = conn.cursor()
                cursor.execute('''
                    INSERT INTO trades (
                        timestamp, market_title, market_slug, model_name,
                        prediction, probability, action, price, amount, size,
                        status, error_message, raw_data
                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                ''', (
                    trade_data.get('timestamp', datetime.now().isoformat()),
                    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('error_message'),
                    json.dumps(trade_data.get('raw_data', {}), ensure_ascii=False)
                ))
                conn.commit()
        except Exception as e:
            print(f"写入数据库失败: {e}")

    def check_existing_trade(self, market_slug: str, status: str = 'success') -> bool:
        """
        检查指定市场是否已有成功交易记录
        
        Args:
            market_slug: 市场slug
            status: 交易状态 (默认 'success', 也可以是 'test')
            
        Returns:
            True 如果已有成功交易, 否则 False
        """
        try:
            with self.get_connection() as conn:
                cursor = conn.cursor()
                cursor.execute('''
                    SELECT COUNT(*) FROM trades 
                    WHERE market_slug = ? AND status = ?
                ''', (market_slug, status))
                count = cursor.fetchone()[0]
                return count > 0
        except Exception as e:
            print(f"检查已存在交易失败: {e}")
            return False