#!/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())
    """
    # 导入所有模型注册到 Base.metadata（必须在 create_all 之前）
    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
    from app.models import scheduled_task  # noqa: F401

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

    # 存量数据回填：将现有任务的最新执行结果转为执行记录（幂等，仅执行一次）
    try:
        from app.services.performance_service import PerformanceService
        async with async_session_maker() as session:
            service = PerformanceService(session)
            migrated = await service._migrate_legacy_executions()
            await session.commit()
            if migrated > 0:
                logger.info(f"存量执行数据回填完成：迁移 {migrated} 条执行记录")
    except Exception as e:
        # 回填失败不阻断启动，下次启动会重试（幂等设计）
        logger.warning(f"存量执行数据回填失败（不影响启动，下次重试）: {e}")


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
    from app.models import scheduled_task  # 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"),
        # 设备模拟：环境配置默认主题参数（含 token 授权码等，旧库升级）
        ("device_env_configs", "default_topic_params", "JSON"),
        # 性能测试：登录接口压测需签名（旧库升级）
        ("performance_tasks", "sign_request", "BOOLEAN DEFAULT 0"),
        # 性能测试：响应捕获规则（旧库升级）
        ("performance_tasks", "capture_rules", "JSON"),
        # 性能测试：任务关联的接口预设（旧库升级）
        ("performance_tasks", "preset_id", "VARCHAR(64) DEFAULT NULL"),
        # 性能测试：唯一性字段配置（旧库升级）
        ("performance_tasks", "unique_fields", "JSON"),
        # 功能测试报告模板：描述字段（旧库升级）
        ("report_templates", "description", "VARCHAR(200) DEFAULT ''"),
        # 钉钉配置：被测系统URL（旧库升级）
        ("dingtalk_configs", "target_url", "VARCHAR(500) DEFAULT ''"),
        # 钉钉配置：登录接口路径（旧库升级）
        ("dingtalk_configs", "login_path", "VARCHAR(200) DEFAULT '/platform/api/auth/login'"),
        # 钉钉配置：验证码接口路径（旧库升级）
        ("dingtalk_configs", "captcha_path", "VARCHAR(200) DEFAULT '/platform/api/code'"),
        # 性能测试：自定义登录用户名（旧库升级）
        ("performance_tasks", "login_username", "VARCHAR(200) DEFAULT NULL"),
        # 性能测试：自定义登录密码（旧库升级）
        ("performance_tasks", "login_password", "VARCHAR(200) DEFAULT NULL"),
        # 定时任务：用例类型（旧库升级，区分 UI/安全测试）
        ("scheduled_tasks", "case_type", "VARCHAR(20) DEFAULT 'ui'"),
        # 定时任务：安全测试配置ID（旧库升级）
        ("scheduled_tasks", "config_id", "VARCHAR(64) DEFAULT NULL"),
        # 性能测试：增强指标体系（P95/标准差/延迟/连接时间/带宽/峰值TPS/Apdex/错误分类，旧库升级）
        ("performance_tasks", "p95_response_time", "FLOAT DEFAULT 0.0"),
        ("performance_tasks", "std_dev", "FLOAT DEFAULT 0.0"),
        ("performance_tasks", "latency_avg", "FLOAT DEFAULT 0.0"),
        ("performance_tasks", "latency_min", "FLOAT DEFAULT 0.0"),
        ("performance_tasks", "latency_max", "FLOAT DEFAULT 0.0"),
        ("performance_tasks", "connect_time_avg", "FLOAT DEFAULT 0.0"),
        ("performance_tasks", "connect_time_max", "FLOAT DEFAULT 0.0"),
        ("performance_tasks", "peak_tps", "FLOAT DEFAULT 0.0"),
        ("performance_tasks", "total_sent_bytes", "INTEGER DEFAULT 0"),
        ("performance_tasks", "total_received_bytes", "INTEGER DEFAULT 0"),
        ("performance_tasks", "apdex", "FLOAT DEFAULT 1.0"),
        ("performance_tasks", "error_type_timeout", "INTEGER DEFAULT 0"),
        ("performance_tasks", "error_type_connect_error", "INTEGER DEFAULT 0"),
        ("performance_tasks", "error_type_client_error", "INTEGER DEFAULT 0"),
        ("performance_tasks", "error_type_http_4xx", "INTEGER DEFAULT 0"),
        ("performance_tasks", "error_type_http_5xx", "INTEGER DEFAULT 0"),
        ("performance_tasks", "error_type_assertion", "INTEGER DEFAULT 0"),
        # 性能测试：快照增强指标（旧库升级）
        ("performance_snapshots", "p95_response_time", "FLOAT DEFAULT 0.0"),
        ("performance_snapshots", "std_dev", "FLOAT DEFAULT 0.0"),
        ("performance_snapshots", "latency_avg", "FLOAT DEFAULT 0.0"),
        ("performance_snapshots", "connect_time_avg", "FLOAT DEFAULT 0.0"),
        ("performance_snapshots", "sent_bytes", "INTEGER DEFAULT 0"),
        ("performance_snapshots", "received_bytes", "INTEGER DEFAULT 0"),
        ("performance_snapshots", "peak_tps", "FLOAT DEFAULT 0.0"),
        # 预设配置：登录凭据名称/用户名/密码（旧库升级）
        ("api_presets", "credential_name", "VARCHAR(100) DEFAULT NULL"),
        ("api_presets", "login_username", "VARCHAR(200) DEFAULT NULL"),
        ("api_presets", "login_password", "VARCHAR(200) DEFAULT NULL"),
        # 性能测试：任务归属项目ID（旧库升级）
        ("performance_tasks", "project_id", "VARCHAR(64) DEFAULT NULL"),
        # 性能测试：HTTP 状态码统计（旧库升级，合并报告聚合用）
        ("performance_tasks", "status_2xx", "INTEGER DEFAULT 0"),
        ("performance_tasks", "status_3xx", "INTEGER DEFAULT 0"),
        ("performance_tasks", "status_4xx", "INTEGER DEFAULT 0"),
        ("performance_tasks", "status_5xx", "INTEGER DEFAULT 0"),
        # 性能测试：事务处理时间（多步骤事务模式，旧库升级）
        ("performance_tasks", "steps", "JSON"),
        ("performance_tasks", "transaction_fail_policy", "VARCHAR(20) DEFAULT 'fail_fast'"),
        ("performance_tasks", "transaction_timeout", "INTEGER DEFAULT 60"),
        ("performance_tasks", "transaction_summary", "JSON"),
        # 性能测试：执行机资源监控汇总（旧库升级）
        ("performance_tasks", "resource_summary", "JSON"),
        # 性能测试：目标机资源监控汇总（SSH 方式采集被测系统资源，旧库升级）
        ("performance_tasks", "target_resource_summary", "JSON"),
        # 性能测试：场景增强字段（瞬时并发/长稳压测，2026-08-22 新增）
        ("performance_tasks", "scenario_type", "VARCHAR(20) DEFAULT NULL"),
        ("performance_tasks", "scenario_apis", "JSON"),
        ("performance_tasks", "loop_type", "VARCHAR(10) DEFAULT 'finite'"),
        ("performance_tasks", "loop_count", "INTEGER DEFAULT 1"),
        ("performance_tasks", "synchronizing_timer_enabled", "BOOLEAN DEFAULT 0"),
        ("performance_tasks", "synchronizing_timer_count", "INTEGER DEFAULT 0"),
        ("performance_tasks", "synchronizing_timer_timeout", "INTEGER DEFAULT 30000"),
        ("performance_tasks", "error_action", "VARCHAR(10) DEFAULT 'continue'"),
        ("performance_tasks", "think_time_enabled", "BOOLEAN DEFAULT 0"),
        ("performance_tasks", "think_time_min", "INTEGER DEFAULT 1000"),
        ("performance_tasks", "think_time_max", "INTEGER DEFAULT 3000"),
        ("performance_tasks", "think_time_distribution", "VARCHAR(10) DEFAULT 'constant'"),
        ("performance_tasks", "csv_parameterization_enabled", "BOOLEAN DEFAULT 0"),
        ("performance_tasks", "csv_content", "TEXT DEFAULT ''"),
        ("performance_tasks", "csv_variable_mapping", "JSON"),
        ("performance_tasks", "api_summary", "JSON"),
        # 性能测试：快照增强指标（快照级 p95 等已于上条添加，本行仅作记录）
        # 性能测试：快照关联执行记录 ID（旧库升级，执行跳转闭环用）
        ("performance_snapshots", "execution_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("数据库表结构已删除")