"""
混合数据库管理器 - 同时支持 Tortoise 和 SQLAlchemy

Version: 0.3.0
Author: Maxwell
"""

from typing import Dict, Optional, List, Any, Literal
from fastapi import FastAPI
from loguru import logger

from infoman.config import settings, DatabaseConfig

# ==================== 类型定义 ====================

ORMType = Literal["tortoise", "sqlalchemy"]


# ==================== 后端实现导入 ====================

# Tortoise 后端
try:
    from tortoise import Tortoise
    from infoman.service.infrastructure.db_relation.mysql import MySQLBackend
    from infoman.service.infrastructure.db_relation.sqllite import SQLiteBackend
    from infoman.service.infrastructure.db_relation.pgsql import PostgreSQLBackend

    TORTOISE_BACKENDS = {
        "mysql": MySQLBackend,
        "postgresql": PostgreSQLBackend,
        "sqlite": SQLiteBackend,
    }
    _TORTOISE_AVAILABLE = True
except ImportError:
    _TORTOISE_AVAILABLE = False
    TORTOISE_BACKENDS = {}

# SQLAlchemy 后端
try:
    from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
    from sqlalchemy import text

    _SQLALCHEMY_AVAILABLE = True
except ImportError:
    _SQLALCHEMY_AVAILABLE = False


# ==================== SQLAlchemy 后端实现 ====================

class SQLAlchemyMySQLBackend:
    """SQLAlchemy MySQL 后端"""

    @staticmethod
    def get_url(config: DatabaseConfig) -> str:
        """生成 MySQL 连接 URL"""
        return (
            f"mysql+asyncmy://{config.user}:{config.password}@"
            f"{config.host}:{config.port}/{config.database}"
            f"?charset={config.charset}"
        )


class SQLAlchemyPostgreSQLBackend:
    """SQLAlchemy PostgreSQL 后端"""

    @staticmethod
    def get_url(config: DatabaseConfig) -> str:
        """生成 PostgreSQL 连接 URL"""
        return (
            f"postgresql+asyncpg://{config.user}:{config.password}@"
            f"{config.host}:{config.port}/{config.database}"
        )


class SQLAlchemySQLiteBackend:
    """SQLAlchemy SQLite 后端"""

    @staticmethod
    def get_url(config: DatabaseConfig) -> str:
        """生成 SQLite 连接 URL"""
        return f"sqlite+aiosqlite:///{config.database}"


SQLALCHEMY_BACKENDS = {
    "mysql": SQLAlchemyMySQLBackend,
    "postgresql": SQLAlchemyPostgreSQLBackend,
    "sqlite": SQLAlchemySQLiteBackend,
}


# ==================== 混合数据库管理器 ====================

class HybridDatabaseManager:
    """
    混合数据库管理器 - 同时支持 Tortoise 和 SQLAlchemy

    功能：
    - 支持多个数据库同时连接
    - 每个连接可独立选择 ORM（Tortoise 或 SQLAlchemy）
    - 统一的健康检查
    - 优雅关闭

    使用示例：
        >>> manager = HybridDatabaseManager()
        >>> await manager.startup(app)
        >>>
        >>> # 获取 Tortoise 连接
        >>> tortoise_conn = manager.get_tortoise_connection("default")
        >>>
        >>> # 获取 SQLAlchemy session maker
        >>> session_maker = manager.get_sqlalchemy_session_maker("analytics")
    """

    def __init__(self):
        # Tortoise 连接
        self.tortoise_connections: Dict[str, DatabaseConfig] = {}
        self.tortoise_initialized = False

        # SQLAlchemy 引擎
        self.sqlalchemy_engines: Dict[str, Any] = {}
        self.sqlalchemy_sessions: Dict[str, Any] = {}
        self.sqlalchemy_initialized = False

    @property
    def is_available(self) -> bool:
        """是否有可用的数据库连接"""
        return (
            (self.tortoise_initialized and len(self.tortoise_connections) > 0) or
            (self.sqlalchemy_initialized and len(self.sqlalchemy_engines) > 0)
        )

    @property
    def tortoise_client(self):
        """获取 Tortoise connections 对象"""
        if not self.tortoise_initialized:
            return None
        if not _TORTOISE_AVAILABLE:
            return None
        return Tortoise.connections

    def get_tortoise_connection(self, name: str = "default"):
        """
        获取 Tortoise 连接

        Args:
            name: 连接名称

        Returns:
            Tortoise connection 对象

        Raises:
            RuntimeError: Tortoise 未初始化
        """
        if not self.tortoise_initialized:
            raise RuntimeError("Tortoise ORM 未初始化")
        if not _TORTOISE_AVAILABLE:
            raise RuntimeError("Tortoise ORM 未安装")
        return Tortoise.connections.get(name)

    def get_sqlalchemy_session_maker(self, name: str = "default"):
        """
        获取 SQLAlchemy Session Maker

        Args:
            name: 连接名称

        Returns:
            async_sessionmaker 实例

        Raises:
            RuntimeError: SQLAlchemy 未初始化
        """
        if not self.sqlalchemy_initialized:
            raise RuntimeError("SQLAlchemy 未初始化")
        return self.sqlalchemy_sessions.get(name)

    def get_sqlalchemy_engine(self, name: str = "default"):
        """获取 SQLAlchemy Engine"""
        if not self.sqlalchemy_initialized:
            raise RuntimeError("SQLAlchemy 未初始化")
        return self.sqlalchemy_engines.get(name)

    def _classify_connections(self) -> tuple[Dict, Dict]:
        """
        分类数据库连接配置

        Returns:
            (tortoise_configs, sqlalchemy_configs)
        """
        enabled_dbs = settings.enabled_databases

        tortoise_configs = {}
        sqlalchemy_configs = {}

        for conn_name, db_config in enabled_dbs.items():
            # 根据配置或命名规则判断 ORM 类型
            orm_type = getattr(db_config, 'orm', None)

            # 如果未指定，使用命名规则判断
            if orm_type is None:
                if conn_name.endswith('_alchemy') or conn_name.endswith('_sqlalchemy'):
                    orm_type = "sqlalchemy"
                else:
                    orm_type = "tortoise"  # 默认使用 Tortoise（向前兼容）

            if orm_type == "tortoise":
                tortoise_configs[conn_name] = db_config
            elif orm_type == "sqlalchemy":
                sqlalchemy_configs[conn_name] = db_config
            else:
                logger.warning(f"⚠️  未知的 ORM 类型: {orm_type}，跳过连接 {conn_name}")

        return tortoise_configs, sqlalchemy_configs

    def _gen_tortoise_config(self, configs: Dict[str, DatabaseConfig]) -> Optional[dict]:
        """生成 Tortoise 配置"""
        if not configs:
            return None

        if not _TORTOISE_AVAILABLE:
            logger.warning("⚠️  Tortoise ORM 未安装，跳过 Tortoise 连接")
            return None

        connections = {}

        for conn_name, db_config in configs.items():
            backend_class = TORTOISE_BACKENDS.get(db_config.type)

            if not backend_class:
                logger.error(f"❌ 不支持的数据库类型: {db_config.type}")
                continue

            connections[conn_name] = {
                "engine": backend_class.get_engine(),
                "credentials": backend_class.get_credentials(db_config),
            }

            self.tortoise_connections[conn_name] = db_config

        if not connections:
            return None

        # 应用配置（模型分组）
        apps = {}
        for conn_name, db_config in configs.items():
            if db_config.models:
                app_name = f"{conn_name}_models"
                apps[app_name] = {
                    "models": db_config.models,
                    "default_connection": conn_name,
                }

        if not apps:
            return None

        config = {
            "connections": connections,
            "apps": apps,
            "use_tz": settings.DB_USE_TZ,
            "timezone": settings.DB_TIMEZONE,
        }

        return config

    async def _init_tortoise(self, app: Optional[FastAPI], configs: Dict[str, DatabaseConfig]):
        """初始化 Tortoise ORM"""
        if not configs:
            return

        config = self._gen_tortoise_config(configs)
        if not config:
            return

        try:
            logger.info("🚀 初始化 Tortoise ORM...")
            for conn_name, db_config in configs.items():
                logger.info(
                    f"   - [{conn_name}] {db_config.type.upper()}: "
                    f"{db_config.user}@{db_config.host}:{db_config.port}/{db_config.database}"
                )

            await Tortoise.init(config=config)

            if app:
                app.state.db_tortoise = self.tortoise_client
                logger.debug("✅ Tortoise 客户端已挂载到 app.state")

            self.tortoise_initialized = True
            logger.success(
                f"✅ Tortoise ORM 连接成功（{len(configs)} 个）\n"
                f"   连接名: {list(configs.keys())}"
            )

        except Exception as e:
            logger.error(f"❌ Tortoise ORM 连接失败: {e}")
            raise

    async def _init_sqlalchemy(self, app: Optional[FastAPI], configs: Dict[str, DatabaseConfig]):
        """初始化 SQLAlchemy"""
        if not configs:
            return

        if not _SQLALCHEMY_AVAILABLE:
            logger.warning("⚠️  SQLAlchemy 未安装，跳过 SQLAlchemy 连接")
            return

        try:
            logger.info("🚀 初始化 SQLAlchemy...")

            for conn_name, db_config in configs.items():
                backend_class = SQLALCHEMY_BACKENDS.get(db_config.type)

                if not backend_class:
                    logger.error(f"❌ 不支持的数据库类型: {db_config.type}")
                    continue

                url = backend_class.get_url(db_config)

                logger.info(
                    f"   - [{conn_name}] {db_config.type.upper()}: "
                    f"{db_config.user}@{db_config.host}:{db_config.port}/{db_config.database}"
                )

                engine = create_async_engine(
                    url,
                    pool_size=db_config.pool_max_size,
                    max_overflow=db_config.pool_max_size,
                    pool_recycle=db_config.pool_recycle,
                    echo=db_config.echo,
                    pool_pre_ping=True,  # 健康检查
                )

                session_maker = async_sessionmaker(
                    engine,
                    class_=AsyncSession,
                    expire_on_commit=False,
                )

                self.sqlalchemy_engines[conn_name] = engine
                self.sqlalchemy_sessions[conn_name] = session_maker

            if self.sqlalchemy_engines:
                if app:
                    app.state.db_sqlalchemy_engines = self.sqlalchemy_engines
                    app.state.db_sqlalchemy_sessions = self.sqlalchemy_sessions
                    logger.debug("✅ SQLAlchemy 引擎已挂载到 app.state")

                self.sqlalchemy_initialized = True
                logger.success(
                    f"✅ SQLAlchemy 连接成功（{len(self.sqlalchemy_engines)} 个）\n"
                    f"   连接名: {list(self.sqlalchemy_engines.keys())}"
                )

        except Exception as e:
            logger.error(f"❌ SQLAlchemy 连接失败: {e}")
            raise

    async def startup(self, app: Optional[FastAPI] = None) -> bool:
        """
        启动数据库管理器

        Args:
            app: FastAPI 应用实例（可选）

        Returns:
            是否成功初始化
        """
        if self.tortoise_initialized or self.sqlalchemy_initialized:
            logger.warning("⚠️  DatabaseManager 已初始化，跳过重复初始化")
            return True

        # 分类连接配置
        tortoise_configs, sqlalchemy_configs = self._classify_connections()

        if not tortoise_configs and not sqlalchemy_configs:
            logger.info("⏭️  数据库未配置，跳过初始化")
            return False

        # 初始化 Tortoise
        if tortoise_configs:
            await self._init_tortoise(app, tortoise_configs)

        # 初始化 SQLAlchemy
        if sqlalchemy_configs:
            await self._init_sqlalchemy(app, sqlalchemy_configs)

        return True

    async def register(self, app: FastAPI) -> bool:
        """注册到 FastAPI 应用"""
        return await self.startup(app)

    async def health_check(self, conn_name: Optional[str] = None) -> Dict:
        """
        健康检查

        Args:
            conn_name: 连接名称（可选，默认检查所有）

        Returns:
            健康检查结果
        """
        if not self.is_available:
            return {
                "status": "not_configured",
                "name": "database",
                "details": {"enabled": False}
            }

        results = {}

        # 检查 Tortoise 连接
        if self.tortoise_initialized:
            for name in self.tortoise_connections.keys():
                if conn_name is None or conn_name == name:
                    results[f"tortoise_{name}"] = await self._check_tortoise_connection(name)

        # 检查 SQLAlchemy 连接
        if self.sqlalchemy_initialized:
            for name in self.sqlalchemy_engines.keys():
                if conn_name is None or conn_name == name:
                    results[f"sqlalchemy_{name}"] = await self._check_sqlalchemy_connection(name)

        # 汇总状态
        all_healthy = all(r.get("status") == "healthy" for r in results.values())

        return {
            "status": "healthy" if all_healthy else "unhealthy",
            "name": "database",
            "details": {
                "connections": results,
                "count": len(results)
            }
        }

    async def _check_tortoise_connection(self, conn_name: str) -> Dict:
        """检查 Tortoise 连接"""
        try:
            conn = Tortoise.connections.get(conn_name)
            db_config = self.tortoise_connections.get(conn_name)

            if not conn or not db_config:
                return {
                    "status": "not_found",
                    "error": f"连接 '{conn_name}' 不存在",
                }

            # 执行健康检查查询
            await conn.execute_query("SELECT 1")

            # 获取连接池状态
            pool_status = {
                "size": conn._pool.size() if hasattr(conn, "_pool") else "N/A",
                "free": conn._pool.freesize() if hasattr(conn, "_pool") else "N/A",
            }

            return {
                "status": "healthy",
                "orm": "tortoise",
                "type": db_config.type,
                "database": db_config.database,
                "pool": pool_status,
            }

        except Exception as e:
            logger.error(f"❌ Tortoise 连接 '{conn_name}' 健康检查失败: {e}")
            return {
                "status": "unhealthy",
                "orm": "tortoise",
                "error": str(e),
            }

    async def _check_sqlalchemy_connection(self, conn_name: str) -> Dict:
        """检查 SQLAlchemy 连接"""
        try:
            engine = self.sqlalchemy_engines.get(conn_name)
            session_maker = self.sqlalchemy_sessions.get(conn_name)

            if not engine or not session_maker:
                return {
                    "status": "not_found",
                    "error": f"连接 '{conn_name}' 不存在",
                }

            # 执行健康检查查询
            async with session_maker() as session:
                result = await session.execute(text("SELECT 1"))
                result.scalar()

            # 获取连接池状态
            pool = engine.pool
            pool_status = {
                "size": pool.size(),
                "checked_in": pool.checkedin(),
                "checked_out": pool.checkedout(),
                "overflow": pool.overflow(),
            }

            return {
                "status": "healthy",
                "orm": "sqlalchemy",
                "type": str(engine.url.drivername),
                "database": engine.url.database,
                "pool": pool_status,
            }

        except Exception as e:
            logger.error(f"❌ SQLAlchemy 连接 '{conn_name}' 健康检查失败: {e}")
            return {
                "status": "unhealthy",
                "orm": "sqlalchemy",
                "error": str(e),
            }

    async def shutdown(self):
        """关闭所有数据库连接"""
        # 关闭 Tortoise
        if self.tortoise_initialized:
            try:
                logger.info("⏹️  关闭 Tortoise ORM 连接...")
                for conn_name in self.tortoise_connections.keys():
                    logger.info(f"   - 关闭连接: {conn_name}")

                await Tortoise.close_connections()
                self.tortoise_initialized = False
                logger.success("✅ Tortoise ORM 连接已关闭")

            except Exception as e:
                logger.error(f"❌ 关闭 Tortoise ORM 连接失败: {e}")

        # 关闭 SQLAlchemy
        if self.sqlalchemy_initialized:
            try:
                logger.info("⏹️  关闭 SQLAlchemy 连接...")
                for conn_name, engine in self.sqlalchemy_engines.items():
                    logger.info(f"   - 关闭引擎: {conn_name}")
                    await engine.dispose()

                self.sqlalchemy_engines.clear()
                self.sqlalchemy_sessions.clear()
                self.sqlalchemy_initialized = False
                logger.success("✅ SQLAlchemy 连接已关闭")

            except Exception as e:
                logger.error(f"❌ 关闭 SQLAlchemy 连接失败: {e}")

    async def close(self):
        """关闭所有数据库连接（兼容旧接口）"""
        await self.shutdown()


# ==================== 全局单例 ====================

hybrid_db_manager = HybridDatabaseManager()


# ==================== 便捷函数 ====================

async def register_databases(app: FastAPI) -> bool:
    """注册数据库到 FastAPI 应用"""
    return await hybrid_db_manager.register(app)


async def check_databases_health(conn_name: Optional[str] = None) -> Dict:
    """检查数据库健康状态"""
    return await hybrid_db_manager.health_check(conn_name)


async def close_databases():
    """关闭所有数据库连接"""
    await hybrid_db_manager.close()


def get_connection_names() -> Dict[str, List[str]]:
    """获取所有连接名称（按 ORM 分类）"""
    return {
        "tortoise": list(hybrid_db_manager.tortoise_connections.keys()),
        "sqlalchemy": list(hybrid_db_manager.sqlalchemy_engines.keys()),
    }
