import pymysql
import pandas as pd
from typing import Optional
from contextlib import contextmanager
import os
from dotenv import load_dotenv
from sqlalchemy import create_engine, text

load_dotenv()


class DatabaseManager:
    """管理 OHLCV 数据的 MySQL 数据库"""

    def __init__(self, db_path: str = None):
        """初始化数据库管理器 (db_path 参数已废弃，仅为兼容性保留)"""
        self._engine = None
        self.initialize_db()
    
    def _get_sqlalchemy_engine(self):
        """获取 SQLAlchemy 引擎（用于 pandas read_sql）"""
        if self._engine is None:
            host = os.getenv('MYSQL_HOST')
            port = os.getenv('MYSQL_PORT', 3306)
            user = os.getenv('MYSQL_USER')
            password = os.getenv('MYSQL_PASSWORD')
            database = os.getenv('MYSQL_DB')
            self._engine = create_engine(
                f"mysql+pymysql://{user}:{password}@{host}:{port}/{database}?charset=utf8mb4",
                pool_pre_ping=True,
                pool_recycle=3600
            )
        return self._engine

    @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',
        )
        try:
            yield conn
        finally:
            conn.close()

    def initialize_db(self):
        """创建表（如果不存在）"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS ohlcv (
                    timestamp DATETIME,
                    symbol VARCHAR(20),
                    timeframe VARCHAR(10),
                    open DOUBLE,
                    high DOUBLE,
                    low DOUBLE,
                    close DOUBLE,
                    volume DOUBLE,
                    PRIMARY KEY (symbol, timeframe, timestamp),
                    INDEX idx_ohlcv_time (symbol, timeframe, timestamp)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
            """)
            conn.commit()

    def save_ohlcv(self, df: pd.DataFrame, symbol: str, timeframe: str):
        """保存 OHLCV 数据到数据库 (Upsert)"""
        if df.empty:
            return

        # 确保 timestamp 是 datetime 类型
        df = df.copy()
        if not pd.api.types.is_datetime64_any_dtype(df['timestamp']):
            df['timestamp'] = pd.to_datetime(df['timestamp'])
            
        # 转换为字符串格式
        df['timestamp'] = df['timestamp'].dt.strftime('%Y-%m-%d %H:%M:%S')
        
        # 添加元数据列
        if 'symbol' not in df.columns:
            df['symbol'] = symbol
        if 'timeframe' not in df.columns:
            df['timeframe'] = timeframe
        
        cols = ['timestamp', 'symbol', 'timeframe', 'open', 'high', 'low', 'close', 'volume']
        data = [tuple(row) for row in df[cols].values]
        
        with self.get_connection() as conn:
            cursor = conn.cursor()
            sql = """
                REPLACE INTO ohlcv 
                (timestamp, symbol, timeframe, open, high, low, close, volume)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
            """
            cursor.executemany(sql, data)
            conn.commit()

    def get_ohlcv(self, symbol: str, timeframe: str, start_date: Optional[str] = None, end_date: Optional[str] = None) -> pd.DataFrame:
        """从数据库获取 OHLCV 数据"""
        
        query = """
            SELECT timestamp, open, high, low, close, volume 
            FROM ohlcv 
            WHERE symbol = :p0 AND timeframe = :p1
        """
        params = {'p0': symbol, 'p1': timeframe}
        param_idx = 2

        if start_date:
            query += f" AND timestamp >= :p{param_idx}"
            params[f'p{param_idx}'] = start_date
            param_idx += 1
        
        if end_date:
            # 处理 'YYYY-MM-DD' 格式
            if len(str(end_date)) == 10:
                query += f" AND timestamp < DATE_ADD(:p{param_idx}, INTERVAL 1 DAY)"
            else:
                query += f" AND timestamp <= :p{param_idx}"
            params[f'p{param_idx}'] = end_date
            
        query += " ORDER BY timestamp ASC"

        engine = self._get_sqlalchemy_engine()
        with engine.connect() as conn:
            df = pd.read_sql_query(text(query), conn, params=params, parse_dates=['timestamp'])
            
        return df

    def get_latest_timestamp(self, symbol: str, timeframe: str) -> Optional[pd.Timestamp]:
        """获取指定 symbol/timeframe 的最新时间戳"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                SELECT MAX(timestamp) 
                FROM ohlcv 
                WHERE symbol = %s AND timeframe = %s
            """, (symbol, timeframe))
            result = cursor.fetchone()
            
        return pd.to_datetime(result[0]) if result and result[0] else None

    def get_earliest_timestamp(self, symbol: str, timeframe: str) -> Optional[pd.Timestamp]:
        """获取指定 symbol/timeframe 的最早时间戳"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                SELECT MIN(timestamp) 
                FROM ohlcv 
                WHERE symbol = %s AND timeframe = %s
            """, (symbol, timeframe))
            result = cursor.fetchone()
            
        return pd.to_datetime(result[0]) if result and result[0] else None