"""异步数据库连接与会话。

连接按当前配置惰性创建，测试可以安全地切换到临时 SQLite 数据库。
"""
from __future__ import annotations

import asyncio
from contextlib import asynccontextmanager
from typing import AsyncIterator, Optional
from weakref import WeakKeyDictionary

from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)

from .config import get_settings
from .models import Base

_engine: Optional[AsyncEngine] = None
_engine_url: Optional[str] = None
_session_factory: Optional[async_sessionmaker[AsyncSession]] = None
_sqlite_session_locks: WeakKeyDictionary[asyncio.AbstractEventLoop, asyncio.Lock] = (
    WeakKeyDictionary()
)


def _sqlite_session_lock() -> asyncio.Lock:
    """Return one transaction lock per event loop for the SQLite prototype.

    SQLite permits only one writer at a time. Serializing the complete
    session_scope transaction prevents the simulator and human control routes
    from reading the same event sequence and then racing on commit. PostgreSQL
    deployments skip this application-level lock.
    """

    loop = asyncio.get_running_loop()
    lock = _sqlite_session_locks.get(loop)
    if lock is None:
        lock = asyncio.Lock()
        _sqlite_session_locks[loop] = lock
    return lock


def get_engine() -> AsyncEngine:
    global _engine, _engine_url, _session_factory
    url = get_settings().database_url
    if _engine is None or _engine_url != url:
        _engine = create_async_engine(url, future=True)
        _engine_url = url
        _session_factory = async_sessionmaker(_engine, expire_on_commit=False)
    return _engine


def get_session_factory() -> async_sessionmaker[AsyncSession]:
    get_engine()
    assert _session_factory is not None
    return _session_factory


@asynccontextmanager
async def _transaction_scope() -> AsyncIterator[AsyncSession]:
    session = get_session_factory()()
    try:
        yield session
        await session.commit()
    except Exception:
        await session.rollback()
        raise
    finally:
        await session.close()


@asynccontextmanager
async def session_scope() -> AsyncIterator[AsyncSession]:
    if get_settings().is_sqlite:
        async with _sqlite_session_lock():
            async with _transaction_scope() as session:
                yield session
        return

    async with _transaction_scope() as session:
        yield session


async def init_db() -> None:
    engine = get_engine()
    async with engine.begin() as connection:
        await connection.run_sync(Base.metadata.create_all)


async def drop_db() -> None:
    engine = get_engine()
    async with engine.begin() as connection:
        await connection.run_sync(Base.metadata.drop_all)


async def dispose_db() -> None:
    global _engine, _engine_url, _session_factory
    if _engine is not None:
        await _engine.dispose()
    _engine = None
    _engine_url = None
    _session_factory = None