"""
数据库 Session 管理
支持 FastAPI/Prefect/CLI 多种场景
"""
import logging
import threading
import time
from contextlib import contextmanager
from functools import wraps
from typing import Generator, Optional, Type, TypeVar, Callable

from sqlalchemy import create_engine, Engine, event, text
from sqlalchemy.exc import OperationalError
from sqlalchemy.orm import sessionmaker, Session, scoped_session
from sqlalchemy.pool import QueuePool

logger = logging.getLogger(__name__)

# 异常重试装饰器
T = TypeVar('T')


def retry_on_deadlock(
        max_retries: int = 3,
        delay: float = 0.1,
        backoff: float = 2.0,
        exceptions: tuple = (OperationalError,)
):
    """
    死锁重试装饰器
    """

    def decorator(func: Callable[..., T]) -> Callable[..., T]:
        @wraps(func)
        def wrapper(*args, **kwargs) -> T:
            mtries, mdelay = max_retries, delay
            while mtries > 0:
                try:
                    return func(*args, **kwargs)
                except exceptions as e:
                    mtries -= 1
                    if mtries == 0:
                        logger.error(f"Max retries exceeded: {e}")
                        raise
                    logger.warning(
                        f"Database error, retrying in {mdelay}s "
                        f"({mtries} retries left): {e}"
                    )
                    time.sleep(mdelay)
                    mdelay *= backoff
            raise RuntimeError("Should not reach here")

        return wrapper

    return decorator


class SessionManager:
    """
    Session 管理器
    支持多种使用场景：
    1. FastAPI 依赖注入
    2. Prefect Task
    3. CLI 脚本
    """

    def __init__(
            self,
            database_uri: str,
            pool_size: int = 10,
            max_overflow: int = 20,
            pool_timeout: int = 30,
            pool_recycle: int = 3600,
            echo: bool = False,
            future: bool = True,
            poolclass: Type = QueuePool,
            isolation_level: str = None,
            connect_args: dict = None,
            **kwargs
    ):
        """
        初始化 Session 管理器

        Args:
            database_uri: 数据库连接 URI
            pool_size: 连接池大小
            max_overflow: 最大溢出连接数
            pool_timeout: 连接超时时间
            pool_recycle: 连接回收时间
            echo: 是否打印 SQL
            future: 使用 SQLAlchemy 2.0 风格
        """
        self.database_uri = database_uri
        self._engine: Optional[Engine] = None
        self._session_factory: Optional[sessionmaker] = None
        self._scoped_session: Optional[scoped_session] = None

        self._metrics = {
            'sessions_created': 0,
            'transactions_committed': 0,
            'transactions_rolled_back': 0,
            'errors': []
        }
        self._lock = threading.RLock()

        # 引擎配置
        self.engine_config = {
            'poolclass': poolclass,
            'pool_size': pool_size,
            'max_overflow': max_overflow,
            'pool_timeout': pool_timeout,
            'pool_recycle': pool_recycle,
            'pool_pre_ping': True,
            'echo': echo,
            'future': future,
            'connect_args': connect_args or {},
            **kwargs
        }

        if isolation_level:
            self.engine_config['isolation_level'] = isolation_level

        # 添加SQLAlchemy事件监听
        self._setup_engine_events()

    def _setup_engine_events(self):
        """设置引擎事件监听"""

        @event.listens_for(Engine, "before_cursor_execute")
        def before_cursor_execute(
                conn, cursor, statement, parameters, context, executemany
        ):
            """SQL执行前"""
            conn.info.setdefault('query_start_time', []).append(time.time())

        @event.listens_for(Engine, "after_cursor_execute")
        def after_cursor_execute(
                conn, cursor, statement, parameters, context, executemany
        ):
            """SQL执行后"""
            total = time.time() - conn.info['query_start_time'].pop()
            if total > 1.0:  # 慢查询阈值
                logger.warning(
                    f"Slow query detected: {total:.3f}s\n"
                    f"SQL: {statement[:200]}..."
                )

    @property
    def engine(self) -> Engine:
        """延迟创建引擎"""
        if self._engine is None:
            self._engine = create_engine(self.database_uri, **self.engine_config)
            logger.info(f"Database engine created: {self._engine.dialect.name}")
        return self._engine

    @property
    def session_factory(self) -> sessionmaker:
        """Session 工厂"""
        if self._session_factory is None:
            self._session_factory = sessionmaker(
                bind=self.engine,
                autocommit=False,
                autoflush=False,
                expire_on_commit=False,
            )
        return self._session_factory

    @property
    def scoped_session_factory(self) -> scoped_session:
        """线程安全的 Scoped Session"""
        if self._scoped_session is None:
            self._scoped_session = scoped_session(self.session_factory)
        return self._scoped_session

    def get_session(self) -> Session:
        """
        获取裸 Session（⚠️ 不自动 commit / rollback）
        仅供框架或高级事务控制使用

        Returns:
            Session 实例
        """
        return self.session_factory()

    @contextmanager
    def session_scope(self) -> Generator[Session, None, None]:
        """
        Session 上下文管理器（用于 Prefect Task 和 CLI）

        使用示例:
            with session_manager.session_scope() as session:
                user = session.query(User).first()
        """
        session = self.session_factory()
        try:
            yield session
            session.commit()
        except Exception as e:
            session.rollback()
            logger.error(f"Session rollback due to: {e}")
            raise
        finally:
            session.close()

    @contextmanager
    def transaction(self) -> Generator[Session, None, None]:
        session = self.session_factory()
        try:
            with session.begin():
                yield session
        except Exception:
            logger.exception("Transaction rollback")
            raise
        finally:
            session.close()

    @contextmanager
    def savepoint(self, session: Session) -> Generator[None, None, None]:
        """
        嵌套事务（SAVEPOINT）

        必须在现有 session 的事务中使用

        使用示例:
            with session_manager.session_scope() as session:
                session.add(parent_data)

                try:
                    with SessionManager.savepoint(session):
                        session.add(child_data)
                        raise Exception("Rollback child only")
                except Exception:
                    pass  # 子事务回滚

                # 父事务继续
        """
        if not session.in_transaction():
            raise RuntimeError("Savepoint requires an active transaction")
        with session.begin_nested():
            yield

    def create_tables(self, base):
        """创建所有表"""
        base.metadata.create_all(bind=self.engine)
        logger.info("Database tables created")

    def drop_tables(self, base):
        """删除所有表"""
        base.metadata.drop_all(bind=self.engine)
        logger.warning("Database tables dropped")

    def close(self):
        """关闭所有连接"""
        if self._scoped_session:
            self._scoped_session.remove()
        if self._engine:
            self._engine.dispose()
            logger.info("Database connections closed")

    def __enter__(self):
        """上下文管理器入口"""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器退出"""
        self.close()


# 全局 Session 管理器实例（可选）
_global_session_manager: Optional[SessionManager] = None


def init_session_manager(database_uri: str, **kwargs) -> SessionManager:
    """
    初始化全局 Session 管理器

    Args:
        database_uri: 数据库连接 URI
        **kwargs: 其他配置参数

    Returns:
        SessionManager 实例
    """
    global _global_session_manager
    _global_session_manager = SessionManager(database_uri, **kwargs)
    return _global_session_manager


def get_session_manager() -> SessionManager:
    """
    获取全局 Session 管理器

    Returns:
        SessionManager 实例

    Raises:
        RuntimeError: 如果未初始化
    """
    if _global_session_manager is None:
        raise RuntimeError(
            "SessionManager not initialized. "
            "Call init_session_manager() first."
        )
    return _global_session_manager


def get_db() -> Generator[Session, None, None]:
    """
    FastAPI 依赖注入：获取数据库会话
    每个请求获得独立的 Session，请求结束后自动关闭
    """
    db = get_session_manager().get_session()
    try:
        yield db
        db.commit()
    except Exception:
        db.rollback()
        raise
    finally:
        db.close()
