#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
模块名称：database.py
模块描述：数据库连接配置和会话管理

作者：czj
创建日期：2026-07-09
最后修改：2026-07-09
"""

import logging
from typing import AsyncGenerator

from sqlalchemy.ext.asyncio import (
    AsyncSession,
    create_async_engine,
    async_sessionmaker
)
from sqlalchemy.orm import DeclarativeBase

from app.config import settings

# 配置日志
logger = logging.getLogger(__name__)

# 创建异步引擎（SQLite 需要特殊配置避免锁定）
# 注意：aiomysql 的 ping() 方法与 SQLAlchemy 的 pool_pre_ping 不兼容
# （AsyncAdapt_aiomysql_connection.ping() 缺少 reconnect 参数），
# 因此 MySQL 改用 pool_recycle 定期回收连接，SQLite 仍用 pool_pre_ping。
_is_sqlite = "sqlite" in settings.DATABASE_URL
engine = create_async_engine(
    settings.DATABASE_URL,
    echo=settings.DEBUG,
    future=True,
    connect_args={"check_same_thread": False} if _is_sqlite else {},
    pool_pre_ping=_is_sqlite,
    pool_recycle=3600 if not _is_sqlite else -1,
)

# 创建异步会话工厂
async_session_maker = async_sessionmaker(
    engine,
    class_=AsyncSession,
    expire_on_commit=False,
    autoflush=False
)


class Base(DeclarativeBase):
    """
    数据库模型基类

    所有数据库模型都应继承此类。

    Example:
        >>> class Module(Base):
        >>>     __tablename__ = "modules"
        >>>     id: Mapped[str] = mapped_column(String(32), primary_key=True)
    """
    pass


async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """
    获取数据库会话（依赖注入）

    用于 FastAPI 路由的依赖注入，自动管理会话生命周期。

    Yields:
        AsyncSession: 数据库会话对象

    Example:
        >>> from fastapi import Depends
        >>>
        >>> @app.get("/modules")
        >>> async def get_modules(db: AsyncSession = Depends(get_db)):
        >>>     result = await db.execute(select(Module))
        >>>     return result.scalars().all()
    """
    async with async_session_maker() as session:
        try:
            yield session
            await session.commit()
        except Exception as e:
            await session.rollback()
            logger.error(f"数据库操作失败: {str(e)}")
            raise
        finally:
            await session.close()


async def init_db() -> None:
    """
    初始化数据库

    创建所有表结构，并对已存在的表补齐新增字段（SQLite 兼容）。
    在应用启动时调用。

    Example:
        >>> import asyncio
        >>> asyncio.run(init_db())
    """
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
        # SQLite 对已存在的表不会自动加列，需手动补齐
        await _ensure_columns(conn)
        logger.info("数据库表结构初始化完成")


async def _ensure_columns(conn) -> None:
    """
    补齐已存在表的新增字段（SQLite ALTER TABLE ADD COLUMN）

    SQLite 的 create_all 只创建不存在的表，已存在的表不会加新列。
    本方法检查并补充新增字段，保证旧库平滑升级。

    Args:
        conn: 异步数据库连接（AsyncConnection），内部通过 run_sync 执行同步检查
    """
    from sqlalchemy import text, inspect

    # 导入模型以确保表被创建（必须导入才会注册到 Base.metadata）
    from app.models import device_sim  # noqa: F401
    from app.models import performance  # noqa: F401
    from app.models import performance_output  # noqa: F401
    from app.models import api_preset  # noqa: F401

    # 需要补齐的字段：[(表名, 列名, 列定义SQL)]
    columns_to_add = [
        ("modules", "module_type", "VARCHAR(20) DEFAULT 'standard'"),
        ("modules", "parent_id", "VARCHAR(32) DEFAULT NULL"),
        ("test_cases", "case_type", "VARCHAR(20) DEFAULT 'ui'"),
        ("executions", "case_type", "VARCHAR(20) DEFAULT 'ui'"),
        # 设备模拟：环境配置多主题字段（旧库升级）
        ("device_env_configs", "topics", "JSON"),
        # 性能测试：登录接口压测需签名（旧库升级）
        ("performance_tasks", "sign_request", "BOOLEAN DEFAULT 0"),
        # 性能测试：响应捕获规则（旧库升级）
        ("performance_tasks", "capture_rules", "JSON"),
        # 性能测试：任务关联的接口预设（旧库升级）
        ("performance_tasks", "preset_id", "VARCHAR(64) DEFAULT NULL"),
    ]

    def _do_ensure(sync_conn) -> None:
        """同步执行：检查并补齐字段（在 run_sync 内调用，可使用 inspect）"""
        inspector = inspect(sync_conn)
        for table, column, col_def in columns_to_add:
            try:
                cols = [c["name"] for c in inspector.get_columns(table)]
                if column not in cols:
                    sync_conn.execute(
                        text(f"ALTER TABLE {table} ADD COLUMN {column} {col_def}")
                    )
                    logger.info(f"补齐字段: {table}.{column}")
            except Exception as e:
                logger.warning(f"补齐字段失败 {table}.{column}: {e}")

    # AsyncConnection 不支持直接 inspect，必须通过 run_sync 在同步上下文中执行
    await conn.run_sync(_do_ensure)


async def drop_db() -> None:
    """
    删除所有数据库表

    用于测试环境清理数据库。
    """
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.drop_all)
        logger.info("数据库表结构已删除")