"""
Apache Doris 数据库优化模块
针对 Doris 特性的优化策略：
1. 连接池配置优化
2. 批量操作优化
3. 查询性能监控
4. 表结构优化建议
"""
import os
import time
import pymysql
from contextlib import contextmanager
from collections import deque
from threading import Lock
from lib.logger import log_info, log_error, log_warning


class DorisConnectionPool:
    """Doris 数据库连接池"""
    
    _instance = None
    _lock = Lock()
    
    def __new__(cls, *args, **kwargs):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._initialized = False
        return cls._instance
    
    def __init__(self, host=None, port=None, user=None, password=None, 
                 database=None, pool_size=10, **kwargs):
        if self._initialized:
            return
            
        self._host = host or os.environ.get('DB_HOST', 'haoslm2.xicp.net')
        self._port = port or int(os.environ.get('DB_PORT', 10216))
        self._user = user or os.environ.get('DB_USER', 'root')
        self._password = password or os.environ.get('DB_PASSWORD', 'DsideaL147258369')
        self._database = database or os.environ.get('DB_NAME', 'yltcharge')
        self._pool_size = pool_size
        
        # 连接池
        self._pool = deque(maxlen=pool_size)
        self._pool_lock = Lock()
        self._total_connections = 0
        self._initialized = True
        
        # 预热连接池
        self._preheat_pool()
        log_info(f"[Doris连接池] 初始化完成，目标连接数: {pool_size}", 'doris')
    
    def _preheat_pool(self):
        """预热连接池"""
        try:
            for i in range(min(3, self._pool_size)):
                try:
                    conn = self._create_connection()
                    with self._pool_lock:
                        self._pool.append(conn)
                        self._total_connections += 1
                    log_info(f"[Doris连接池] 预热连接 {i+1}/3 成功", 'doris')
                except Exception as e:
                    log_warning(f"[Doris连接池] 预热连接失败: {e}", 'doris')
        except Exception as e:
            log_warning(f"[Doris连接池] 预热异常: {e}", 'doris')
    
    def _create_connection(self):
        """创建新连接"""
        conn = pymysql.connect(
            host=self._host,
            port=self._port,
            user=self._user,
            password=self._password,
            database=self._database,
            charset='utf8mb4',
            cursorclass=pymysql.cursors.DictCursor,
            autocommit=True,
            connect_timeout=10,  # 缩短连接超时
            read_timeout=120,
            write_timeout=120,
            # Doris 特有优化
            init_command='SET SESSION query_timeout = 300000',  # 5分钟超时
        )
        return conn
    
    def get_connection(self, timeout=5):
        """获取连接（带超时）"""
        start_time = time.time()
        
        # 先从池中获取
        with self._pool_lock:
            if self._pool:
                conn = self._pool.pop()
                # 检查连接是否有效
                try:
                    conn.ping(reconnect=False)
                    return conn
                except Exception:
                    # 连接失效，创建新连接
                    self._total_connections -= 1
                    try:
                        conn.close()
                    except Exception:
                        pass
        
        # 池没有可用连接，等待或创建新连接
        elapsed = time.time() - start_time
        if elapsed < timeout:
            # 尝试创建新连接（如果总数未超过限制）
            if self._total_connections < self._pool_size:
                try:
                    conn = self._create_connection()
                    with self._pool_lock:
                        self._total_connections += 1
                    return conn
                except Exception as e:
                    log_warning(f"[Doris连接池] 创建新连接失败: {e}", 'doris')
        
        # 超时或达到上限，强制创建新连接
        log_warning(f"[Doris连接池] 强制创建新连接（池已满或超时）", 'doris')
        return self._create_connection()
    
    def return_connection(self, conn):
        """归还连接到池"""
        if conn is None:
            return
        try:
            conn.ping(reconnect=False)
            with self._pool_lock:
                if len(self._pool) < self._pool_size:
                    self._pool.append(conn)
                    return
        except Exception:
            pass
        
        # 连接失效或池已满，关闭连接
        try:
            conn.close()
        except Exception:
            pass
        with self._pool_lock:
            self._total_connections = max(0, self._total_connections - 1)
    
    def get_stats(self):
        """获取连接池状态"""
        with self._pool_lock:
            return {
                'pool_size': len(self._pool),
                'total_connections': self._total_connections,
                'max_pool_size': self._pool_size
            }


# 全局连接池实例
_pool = None


def get_pool():
    """获取全局连接池"""
    global _pool
    if _pool is None:
        _pool = DorisConnectionPool(pool_size=10)
    return _pool


@contextmanager
def get_doris_connection():
    """获取 Doris 连接的上下文管理器"""
    pool = get_pool()
    conn = pool.get_connection(timeout=3)
    try:
        yield conn
    except pymysql.Error as e:
        log_error(f"[Doris] 数据库操作失败: {e}", 'doris')
        # 连接可能已失效，不归还到池
        try:
            conn.close()
        except Exception:
            pass
        with pool._pool_lock:
            pool._total_connections = max(0, pool._total_connections - 1)
        raise
    else:
        pool.return_connection(conn)


def execute_doris_query(sql, params=None, retry=2):
    """
    执行 Doris 查询（带重试和性能监控）
    
    Args:
        sql: SQL 语句
        params: 参数
        retry: 重试次数
    
    Returns:
        list: 查询结果
    """
    for attempt in range(retry + 1):
        try:
            start_time = time.time()
            
            with get_doris_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute(sql, params)
                    result = cursor.fetchall()
            
            elapsed = (time.time() - start_time) * 1000  # 毫秒
            
            # 慢查询监控（超过100ms记录）
            if elapsed > 100:
                log_info(f"[Doris慢查询] 耗时: {elapsed:.1f}ms, SQL: {sql[:200]}, 行数: {len(result) if result else 0}", 'doris')
            
            if attempt > 0:
                log_info(f'[Doris] 查询重试成功，第{attempt+1}次尝试', 'doris')
            
            return result
            
        except pymysql.err.OperationalError as e:
            if attempt < retry and (e.args[0] == 2013 or e.args[0] == 2006):
                log_warning(f'[Doris] 查询连接断开，正在重试（第{attempt+1}次）: {e}', 'doris')
                time.sleep(0.5 * (attempt + 1))
                continue
            log_error(f'[Doris] 查询失败: {e}\nSQL: {sql}\nParams: {params}', 'doris')
            raise
        except Exception as e:
            log_error(f'[Doris] 查询失败: {e}\nSQL: {sql}\nParams: {params}', 'doris')
            raise


def execute_doris_batch_insert(sql, params_list, batch_size=100, retry=2):
    """
    Doris 批量插入优化
    
    Args:
        sql: SQL 语句（带 %s 占位符）
        params_list: 参数列表
        batch_size: 每批大小
        retry: 重试次数
    
    Returns:
        int: 插入的总行数
    """
    if not params_list:
        return 0
    
    total_inserted = 0
    start_time = time.time()
    
    for i in range(0, len(params_list), batch_size):
        batch = params_list[i:i + batch_size]
        batch_sql = sql.rstrip()
        
        # 对于 Doris，可以使用 INSERT INTO ... VALUES (...), (...), (...) 格式
        # 但 pymysql 的 executemany 已经处理好了
        
        for attempt in range(retry + 1):
            try:
                with get_doris_connection() as conn:
                    with conn.cursor() as cursor:
                        cursor.executemany(batch_sql, batch)
                        total_inserted += len(batch)
                
                if attempt > 0:
                    log_info(f'[Doris] 批量插入重试成功，第{attempt+1}次尝试', 'doris')
                break
                
            except pymysql.err.OperationalError as e:
                if attempt < retry and (e.args[0] == 2013 or e.args[0] == 2006):
                    log_warning(f'[Doris] 批量插入连接断开，正在重试（第{attempt+1}次）: {e}', 'doris')
                    time.sleep(0.5 * (attempt + 1))
                    continue
                log_error(f'[Doris] 批量插入失败: {e}\nSQL: {batch_sql}', 'doris')
                raise
            except Exception as e:
                log_error(f'[Doris] 批量插入失败: {e}\nSQL: {batch_sql}', 'doris')
                raise
    
    elapsed = (time.time() - start_time) * 1000
    log_info(f'[Doris] 批量插入完成，共 {total_inserted} 行，耗时 {elapsed:.1f}ms', 'doris')
    
    return total_inserted


def execute_doris_update(sql, params=None, retry=2):
    """
    执行 Doris 更新操作（带重试）
    
    Args:
        sql: SQL 语句
        params: 参数
        retry: 重试次数
    
    Returns:
        int: 影响的行数
    """
    for attempt in range(retry + 1):
        try:
            start_time = time.time()
            
            with get_doris_connection() as conn:
                with conn.cursor() as cursor:
                    affected = cursor.execute(sql, params)
            
            elapsed = (time.time() - start_time) * 1000
            
            if elapsed > 100:
                log_info(f"[Doris慢更新] 耗时: {elapsed:.1f}ms, SQL: {sql[:200]}", 'doris')
            
            if attempt > 0:
                log_info(f'[Doris] 更新重试成功，第{attempt+1}次尝试', 'doris')
            
            return cursor.rowcount if hasattr(cursor, 'rowcount') else affected
            
        except pymysql.err.OperationalError as e:
            if attempt < retry and (e.args[0] == 2013 or e.args[0] == 2006):
                log_warning(f'[Doris] 更新连接断开，正在重试（第{attempt+1}次）: {e}', 'doris')
                time.sleep(0.5 * (attempt + 1))
                continue
            log_error(f'[Doris] 更新失败: {e}\nSQL: {sql}\nParams: {params}', 'doris')
            raise
        except Exception as e:
            log_error(f'[Doris] 更新失败: {e}\nSQL: {sql}\nParams: {params}', 'doris')
            raise


def optimize_doris_table(table_name):
    """
    优化 Doris 表的统计信息
    
    Args:
        table_name: 表名
    """
    try:
        sql = f"ANALYZE TABLE {table_name} UPDATE HISTOGRAM"
        execute_doris_update(sql)
        log_info(f"[Doris] 表 {table_name} 统计信息已更新", 'doris')
        return True
    except Exception as e:
        log_warning(f"[Doris] 表 {table_name} 统计信息更新失败: {e}", 'doris')
        # 尝试使用 COMPUTE STATISTICS
        try:
            sql = f"ANALYZE TABLE {table_name} COMPUTE STATISTICS"
            execute_doris_update(sql)
            log_info(f"[Doris] 表 {table_name} 统计信息已计算", 'doris')
            return True
        except Exception as e2:
            log_warning(f"[Doris] 表 {table_name} 统计信息计算也失败: {e2}", 'doris')
            return False


def get_doris_table_stats(table_name):
    """
    获取 Doris 表的统计信息
    
    Args:
        table_name: 表名
    
    Returns:
        dict: 统计信息
    """
    try:
        sql = f"SHOW TABLE STATUS LIKE '{table_name}'"
        result = execute_doris_query(sql)
        if result:
            return result[0]
        return {}
    except Exception as e:
        log_warning(f"[Doris] 获取表 {table_name} 统计信息失败: {e}", 'doris')
        return {}


def print_pool_stats():
    """打印连接池状态"""
    pool = get_pool()
    stats = pool.get_stats()
    log_info(f"[Doris连接池] 状态: {stats}", 'doris')
    return stats