"""数据库连接池管理 — 复用数据库连接，提高性能。

使用异步连接池:
- MySQL: aiomysql.Pool
- PostgreSQL: asyncpg.Pool

连接池按实例 ID 管理，每个实例维护一个连接池。
"""
import asyncio
import time
from dataclasses import dataclass, field
from typing import Dict, Optional

from app.models.instance import Instance
from app.core.crypto import decrypt_password


@dataclass
class PoolStats:
    """连接池统计信息。"""
    instance_id: int
    db_type: str
    active_connections: int = 0
    idle_connections: int = 0
    max_connections: int = 5
    total_queries: int = 0
    failed_queries: int = 0
    avg_response_ms: int = 0
    pool_hit_rate: float = 1.0
    created_at: float = field(default_factory=time.time)

    def to_dict(self) -> dict:
        return {
            "instance_id": self.instance_id,
            "db_type": self.db_type,
            "active_connections": self.active_connections,
            "idle_connections": self.idle_connections,
            "max_connections": self.max_connections,
            "total_queries": self.total_queries,
            "failed_queries": self.failed_queries,
            "avg_response_ms": self.avg_response_ms,
            "pool_hit_rate": round(self.pool_hit_rate, 4),
        }


class ConnectionPoolManager:
    """数据库连接池管理器。"""

    def __init__(self):
        # 连接池缓存: {instance_id: pool}
        self._pools: Dict[int, Any] = {}
        # 连接池统计
        self._stats: Dict[int, PoolStats] = {}
        # 创建锁，防止并发创建重复连接池
        self._locks: Dict[int, asyncio.Lock] = {}
        # 总体锁
        self._global_lock = asyncio.Lock()

    async def get_pool(self, instance: Instance):
        """获取或创建连接池。"""
        instance_id = instance.id

        # 检查现有连接池
        if instance_id in self._pools:
            pool = self._pools[instance_id]
            # 简单的健康检查
            try:
                if instance.db_type == "mysql":
                    pool._free  # 检查 aiomysql pool 是否活跃
                return pool
            except Exception:
                # 连接池已失效，重新创建
                await self.close_pool(instance_id)

        # 获取或创建实例锁
        if instance_id not in self._locks:
            self._locks[instance_id] = asyncio.Lock()

        async with self._locks[instance_id]:
            # 双重检查
            if instance_id in self._pools:
                return self._pools[instance_id]

            pool = await self._create_pool(instance)
            self._pools[instance_id] = pool
            self._stats[instance_id] = PoolStats(
                instance_id=instance_id,
                db_type=instance.db_type,
                max_connections=instance.connection_pool_size or 5,
            )
            return pool

    async def _create_pool(self, instance: Instance):
        """创建连接池。"""
        password = decrypt_password(instance.password)
        pool_size = instance.connection_pool_size or 5

        if instance.db_type == "mysql":
            import aiomysql
            return await aiomysql.create_pool(
                host=instance.host,
                port=instance.port,
                user=instance.username,
                password=password,
                db=instance.database,
                minsize=1,
                maxsize=pool_size,
                pool_recycle=3600,       # 1小时回收连接
                connect_timeout=10,
                autocommit=True,
            )
        elif instance.db_type in ("postgres", "postgresql"):
            import asyncpg
            dsn = (
                f"postgresql://{instance.username}:{password}"
                f"@{instance.host}:{instance.port}"
                f"/{instance.database or 'postgres'}"
            )
            return await asyncpg.create_pool(
                dsn=dsn,
                min_size=1,
                max_size=pool_size,
                command_timeout=60,
                max_inactive_connection_lifetime=3600,
            )
        else:
            raise ValueError(f"不支持的连接池类型: {instance.db_type}")

    async def close_pool(self, instance_id: int) -> None:
        """关闭指定实例的连接池。"""
        pool = self._pools.pop(instance_id, None)
        if pool:
            try:
                pool.close()
                await pool.wait_closed()
            except Exception:
                pass
        self._stats.pop(instance_id, None)

    async def close_all(self) -> None:
        """关闭所有连接池。"""
        for instance_id in list(self._pools.keys()):
            await self.close_pool(instance_id)
        self._pools.clear()
        self._stats.clear()

    async def get_connection(self, instance: Instance):
        """从连接池获取一个连接（上下文管理器）。"""
        pool = await self.get_pool(instance)
        return pool.acquire()

    def get_pool_stats(self, instance_id: int) -> Optional[PoolStats]:
        """获取连接池统计。"""
        return self._stats.get(instance_id)

    def get_all_stats(self) -> Dict[int, PoolStats]:
        """获取所有连接池统计。"""
        return dict(self._stats)

    async def health_check(self, instance_id: int) -> bool:
        """检查连接池健康状态。"""
        pool = self._pools.get(instance_id)
        if not pool:
            return False

        try:
            if hasattr(pool, '_free'):
                # aiomysql
                conn = await pool.acquire()
                try:
                    async with conn.cursor() as cur:
                        await cur.execute("SELECT 1")
                        await cur.fetchone()
                    return True
                finally:
                    pool.release(conn)
            else:
                # asyncpg
                conn = await pool.acquire()
                try:
                    await conn.fetchval("SELECT 1")
                    return True
                finally:
                    await pool.release(conn)
        except Exception:
            return False

    async def reset_pool(self, instance_id: int) -> None:
        """重置连接池（关闭并重新创建）。"""
        await self.close_pool(instance_id)

    def update_stats(self, instance_id: int, query_success: bool = True,
                     response_ms: int = 0) -> None:
        """更新查询统计。"""
        stats = self._stats.get(instance_id)
        if not stats:
            return

        stats.total_queries += 1
        if not query_success:
            stats.failed_queries += 1

        # 移动平均响应时间
        if response_ms > 0:
            if stats.avg_response_ms == 0:
                stats.avg_response_ms = response_ms
            else:
                stats.avg_response_ms = (stats.avg_response_ms * 9 + response_ms) // 10


# 全局连接池管理器实例
_pool_manager: Optional[ConnectionPoolManager] = None


def get_pool_manager() -> ConnectionPoolManager:
    """获取全局连接池管理器。"""
    global _pool_manager
    if _pool_manager is None:
        _pool_manager = ConnectionPoolManager()
    return _pool_manager


async def close_all_pools():
    """关闭所有连接池（应用关闭时调用）。"""
    global _pool_manager
    if _pool_manager:
        await _pool_manager.close_all()
        _pool_manager = None