# _*_ coding : UTF-8 _*_
# @Time : 2025/01/18 02:00
# @UpdateTime : 2025/01/18 02:00
# @Author : sonder
# @File : database.py
# @Software : PyCharm
# @Comment : 本程序
import logging
import sys
from logging.handlers import RotatingFileHandler

from tortoise import Tortoise

from config.env import DataBaseConfig
from utils.log import logger, log_path_sql


async def init_db():
    """
    异步初始化数据库连接。
    """
    # 在数据库连接 URL 中添加时区参数（东八区）
    db_url = (
        f"mysql://{DataBaseConfig.db_username}:{DataBaseConfig.db_password}@"
        f"{DataBaseConfig.db_host}:{DataBaseConfig.db_port}/{DataBaseConfig.db_database}"
        "?charset=utf8mb4"  # 指定时区为东八区,
    )

    await Tortoise.init(
        db_url=db_url,
        modules={"models": ["models"]},  # 指向 models 目录,
        timezone="Asia/Shanghai",
    )

    # 根据 db_echo 配置是否打印 SQL 查询日志
    if DataBaseConfig.db_echo:
        logger.info("SQL 查询日志已启用")
        await configure_tortoise_logging(enable_logging=True, log_level=DataBaseConfig.db_log_level)
    else:
        logger.info("SQL 查询日志已禁用")
        # 禁用 SQL 查询日志
        logger.remove(log_path_sql)

    # 生成数据库表结构
    await Tortoise.generate_schemas()
    logger.success("数据库连接成功！")


async def close_db():
    """
    关闭数据库连接。
    """
    await Tortoise.close_connections()
    logger.success("数据库连接关闭！")


async def configure_tortoise_logging(enable_logging: bool = True, log_level: int = logging.DEBUG):
    """
    异步配置 Tortoise ORM 日志输出。

    :param enable_logging: 是否启用日志输出
    :param log_level: 日志输出级别，默认为 DEBUG
    """
    aiomysql_logger = logging.getLogger("aiomysql")
    tortoise_logger = logging.getLogger("tortoise")

    # 清除之前的处理器，避免重复添加
    if tortoise_logger.hasHandlers():
        tortoise_logger.handlers.clear()

    if aiomysql_logger.hasHandlers():
        aiomysql_logger.handlers.clear()

    if enable_logging:
        # 设置日志格式
        fmt = logging.Formatter(
            fmt="%(asctime)s - %(name)s:%(lineno)d - %(levelname)s - %(message)s",
            datefmt="%Y-%m-%d %H:%M:%S",
        )

        # 创建控制台处理器（输出到控制台）
        console_handler = logging.StreamHandler(sys.stdout)
        console_handler.setLevel(log_level)
        console_handler.setFormatter(fmt)

        # 创建文件处理器（输出到文件）
        file_handler = RotatingFileHandler(
            filename=log_path_sql,
            maxBytes=50 * 1024 * 1024,  # 日志文件大小达到 50MB 时轮换
            backupCount=5,  # 保留 5 个旧日志文件
            encoding="utf-8",
        )
        file_handler.setLevel(log_level)
        file_handler.setFormatter(fmt)

        # 配置 tortoise 顶级日志记录器
        tortoise_logger.setLevel(log_level)
        tortoise_logger.addHandler(console_handler)  # 添加控制台处理器
        tortoise_logger.addHandler(file_handler)  # 添加文件处理器

        # 配置 aiomysql 日志记录器
        aiomysql_logger.setLevel(log_level)
        aiomysql_logger.addHandler(console_handler)  # 添加控制台处理器
        aiomysql_logger.addHandler(file_handler)  # 添加文件处理器
        # 配置 SQL 查询日志记录器
        sql_logger = logging.getLogger("tortoise.db_client")
        sql_logger.setLevel(log_level)

        class SQLResultLogger(logging.Handler):
            async def emit(self, record):
                # 只处理 SQL 查询相关的日志
                if "SELECT" in record.getMessage() or "INSERT" in record.getMessage() or "UPDATE" in record.getMessage() or "DELETE" in record.getMessage():
                    # 输出 SQL 查询语句
                    console_handler.emit(record)
                    file_handler.emit(record)

                    # 异步获取并记录查询结果
                    await self.log_query_result(record)

            async def log_query_result(self, record):
                """
                执行查询并返回结果。
                """
                try:
                    from tortoise import Tortoise
                    connection = Tortoise.get_connection("default")
                    result = await connection.execute_query_dict(record.getMessage())
                    return result
                except Exception as e:
                    return f"获取查询结果失败: {str(e)}"

        # 添加自定义 SQL 查询日志处理器
        sql_result_handler = SQLResultLogger()
        sql_result_handler.setLevel(log_level)
        sql_logger.addHandler(sql_result_handler)
    else:
        # 如果禁用日志，设置日志级别为 WARNING 以抑制大部分输出
        tortoise_logger.setLevel(logging.WARNING)
