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

load_dotenv()

class TradingDatabase:
    """交易数据库管理类 - 用于记录交易历史和预测结果
    
    支持 SQLite (本地) 和 MySQL (远程)。
    如果 .env 中配置了 MYSQL_HOST，则优先使用 MySQL。
    """
    
    def __init__(self, db_path="trading_history.db"):
        self.db_path = db_path
        self.use_mysql = bool(os.getenv('MYSQL_HOST'))
        self._init_db()

    @contextmanager
    def get_connection(self):
        """数据库连接上下文管理器"""
        if self.use_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
            )
        else:
            conn = sqlite3.connect(self.db_path)
            # Make sqlite behave like dict cursor for consistency if needed, 
            # but existing code might expect tuple.
            # Actually existing code uses cursor.execute directly. 
            # SQLite cursor returns tuples by default.
            # PyMySQL DictCursor returns dicts.
            # We need to be careful about compatibility.
            # The existing log_trade uses INSERT which is fine.
            # check_existing_trade uses SELECT count(*), fetchone()[0], which works for both if we don't use DictCursor for MySQL or handle it.
            # Let's use standard cursor for MySQL for compatibility with existing fetchone()[0] usage.
            
        try:
            yield conn
        finally:
            conn.close()

    def _init_db(self):
        """初始化数据库表结构"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            if self.use_mysql:
                # MySQL Schema
                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;
                ''')
            else:
                # SQLite Schema
                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):
        """
        记录交易信息
        """
        try:
            with self.get_connection() as conn:
                cursor = conn.cursor()
                
                # Prepare data
                timestamp = trade_data.get('timestamp', datetime.now().isoformat())
                raw_data_json = json.dumps(trade_data.get('raw_data', {}), ensure_ascii=False)
                
                if self.use_mysql:
                    sql = '''
                        INSERT INTO trades (
                            timestamp, market_title, market_slug, model_name,
                            prediction, probability, action, price, amount, size,
                            status, error_message, raw_data
                        ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                    '''
                else:
                    sql = '''
                        INSERT INTO trades (
                            timestamp, market_title, market_slug, model_name,
                            prediction, probability, action, price, amount, size,
                            status, error_message, raw_data
                        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    '''
                
                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('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()
                
                if self.use_mysql:
                    sql = 'SELECT COUNT(*) FROM trades WHERE market_slug = %s AND status = %s'
                else:
                    sql = 'SELECT COUNT(*) FROM trades WHERE market_slug = ? AND status = ?'
                    
                cursor.execute(sql, (market_slug, status))
                
                # Fetch result compatible with both
                result = cursor.fetchone()
                if isinstance(result, dict): # If DictCursor was used (not used here but for safety)
                    count = list(result.values())[0]
                elif isinstance(result, (list, tuple)):
                    count = result[0]
                else:
                    count = 0
                    
                return count > 0
        except Exception as e:
            print(f"检查已存在交易失败: {e}")
            return False