import logging
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from contextvars import ContextVar
from dataclasses import dataclass
from functools import partial

from sqlalchemy.ext.asyncio import AsyncEngine
from sqlalchemy.ext.asyncio import async_sessionmaker
from sqlalchemy.ext.asyncio import create_async_engine
from sqlmodel.ext.asyncio.session import AsyncSession
from sqlmodel.pool import StaticPool

from ..config import Settings
from ..config import get_settings

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class DatabaseRuntime:
    """Database runtime resources owned by app lifespan."""

    engine: AsyncEngine
    session_factory: async_sessionmaker[AsyncSession]


_runtime: ContextVar[DatabaseRuntime | None] = ContextVar("database_runtime", default=None)


async def dispose_database_runtime() -> None:
    """Dispose lifespan-owned async database resources."""
    runtime = _runtime.get()
    if runtime is None:
        return
    await runtime.engine.dispose()
    _runtime.set(None)


def _to_async_database_url(database_url: str) -> str:
    """Normalize configured database URL to an async SQLAlchemy driver URL."""
    if database_url.startswith("sqlite://") and not database_url.startswith("sqlite+aiosqlite://"):
        return database_url.replace("sqlite://", "sqlite+aiosqlite://", 1)
    if database_url.startswith("postgresql://") and not database_url.startswith("postgresql+asyncpg://"):
        return database_url.replace("postgresql://", "postgresql+asyncpg://", 1)
    return database_url


def _build_engine(settings: Settings) -> AsyncEngine:
    database_url = _to_async_database_url(settings.database_url)
    engine_factory = partial(
        create_async_engine,
        url=database_url,
        echo=False,
        pool_pre_ping=True,
    )

    if database_url.startswith("sqlite"):
        sqlite_connect_settings = {"check_same_thread": settings.sqlite_check_same_thread}
        engine_factory = partial(engine_factory, connect_args=sqlite_connect_settings)
        if ":memory:" in database_url:
            engine_factory = partial(engine_factory, poolclass=StaticPool)

    return engine_factory()


def initialize_database_runtime(*, settings: Settings | None = None) -> DatabaseRuntime:
    """Initialize lifespan-owned async DB resources once per process."""
    runtime = _runtime.get()
    if runtime is not None:
        return runtime

    active_settings = settings or get_settings()
    engine = _build_engine(active_settings)
    session_factory = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
    runtime = DatabaseRuntime(engine=engine, session_factory=session_factory)
    _runtime.set(runtime)
    logger.debug("Initialized async database runtime for database_url=%s", engine.url)
    return runtime


def get_engine(settings: Settings | None = None) -> AsyncEngine:
    """Return the current async SQLAlchemy engine."""
    runtime = _runtime.get() or initialize_database_runtime(settings=settings)
    return runtime.engine


def get_session_factory(settings: Settings | None = None) -> async_sessionmaker[AsyncSession]:
    """Return the shared async session factory."""
    runtime = _runtime.get() or initialize_database_runtime(settings=settings)
    return runtime.session_factory


@asynccontextmanager
async def get_session(
    *,
    settings: Settings | None = None,
    session_factory: async_sessionmaker[AsyncSession] | None = None,
) -> AsyncGenerator[AsyncSession]:
    """Yield a database session and ensure cleanup."""
    active_session_factory = session_factory or get_session_factory(settings)
    async with active_session_factory() as session:
        yield session
