from typing import Optional
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from contextlib import contextmanager
from loguru import logger
from app.config import get_settings


class DatabaseEngine:
    """数据库操作引擎"""

    def __init__(self):
        settings = get_settings()
        self.engine = create_engine(
            settings.DATABASE_URL,
            pool_size=20,
            max_overflow=10,
            pool_pre_ping=True,
            echo=settings.DEBUG
        )
        self.Session = sessionmaker(bind=self.engine)

    @contextmanager
    def get_session(self):
        """获取数据库会话（上下文管理器）"""
        session = self.Session()
        try:
            yield session
            session.commit()
        except Exception as e:
            session.rollback()
            logger.error(f"数据库操作失败：{e}")
            raise
        finally:
            session.close()

    def execute_sql(self, sql: str, params: Optional[dict] = None) -> list:
        """执行单条 SQL"""
        with self.get_session() as session:
            result = session.execute(text(sql), params or {})
            return result.fetchall()

    def execute_transaction(self, sql_list: list, params_list: Optional[list] = None) -> bool:
        """执行事务（多条 SQL）"""
        params_list = params_list or [None] * len(sql_list)
        with self.get_session() as session:
            for sql, params in zip(sql_list, params_list):
                session.execute(text(sql), params or {})
        return True

    def get_table_structure(self, table_name: str):
        """获取表结构（安全参数化查询）"""
        sql = """
        SELECT
            COLUMN_NAME,
            DATA_TYPE,
            CHARACTER_MAXIMUM_LENGTH,
            IS_NULLABLE,
            COLUMN_DEFAULT
        FROM INFORMATION_SCHEMA.COLUMNS
        WHERE TABLE_NAME = :table_name
        ORDER BY ORDINAL_POSITION
        """
        return self.execute_sql(sql, {"table_name": table_name})

    def table_exists(self, table_name: str) -> bool:
        """检查表是否存在（安全参数化查询）"""
        sql = """
        SELECT COUNT(*)
        FROM INFORMATION_SCHEMA.TABLES
        WHERE TABLE_NAME = :table_name
        """
        result = self.execute_sql(sql, {"table_name": table_name})
        return result[0][0] > 0

    def dispose(self):
        """关闭连接池，释放资源"""
        self.engine.dispose()