import threading

from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine

from app.core.config import settings

_engine = None
_session_factory: async_sessionmaker[AsyncSession] | None = None
_lock = threading.RLock()


def _get_engine():
    global _engine
    if _engine is None:
        with _lock:
            if _engine is None:
                _engine = create_async_engine(
                    settings.database_url,
                    echo=settings.debug,
                    pool_pre_ping=True,
                    pool_recycle=300,
                )
    return _engine


def _get_session_factory():
    global _session_factory
    if _session_factory is None:
        with _lock:
            if _session_factory is None:
                _session_factory = async_sessionmaker(
                    _get_engine(), class_=AsyncSession, expire_on_commit=False
                )
    return _session_factory


async def dispose_engine():
    """Dispose the async engine and reset cached DB handles.

    Celery tasks run async code through ``asyncio.run()``, which creates and closes
    one event loop per task. Asyncpg connections are bound to the loop that created
    them, so worker tasks must not keep pooled connections across task runs.
    """
    global _engine, _session_factory

    engine = _engine
    if engine is not None:
        await engine.dispose()

    with _lock:
        if _engine is engine:
            _engine = None
            _session_factory = None


async def init_db():
    """Create tables if they do not exist."""
    from app.db.models import Base  # main models

    engine = _get_engine()
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)


async def get_db():
    """Yield a DB session with proper cleanup."""
    session_factory = _get_session_factory()
    async with session_factory() as session:
        yield session