from datetime import datetime
import os

from sqlalchemy import DateTime, func
from sqlalchemy.ext.asyncio import AsyncAttrs, async_sessionmaker, create_async_engine
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from sqlalchemy.pool import NullPool

from app.config import settings


def _is_celery_process() -> bool:
    argv = " ".join(os.sys.argv).lower()
    return "celery" in argv


engine_kwargs = {
    "echo": settings.DEBUG,
    "pool_pre_ping": True,
}

# SQLite 本地调试时不要乱塞 pool_size/max_overflow，PostgreSQL/asyncpg 才建议配置。
# 默认保持 Celery 连接池复用；只有显式开启 CELERY_DB_USE_NULLPOOL=true 时才降级 NullPool。
if _is_celery_process() and not settings.DATABASE_URL.startswith("sqlite"):
    if bool(getattr(settings, "CELERY_DB_USE_NULLPOOL", False)):
        engine_kwargs.update(poolclass=NullPool)
    else:
        engine_kwargs.update(
            pool_size=settings.CELERY_DB_POOL_SIZE,
            max_overflow=settings.CELERY_DB_MAX_OVERFLOW,
            pool_timeout=settings.CELERY_DB_POOL_TIMEOUT,
            pool_recycle=settings.CELERY_DB_POOL_RECYCLE,
        )

engine = create_async_engine(settings.DATABASE_URL, **engine_kwargs)

async_session = async_sessionmaker(
    engine,
    expire_on_commit=False,
    autoflush=False,
)

AsyncSessionLocal = async_session


class Base(AsyncAttrs, DeclarativeBase):
    pass


class TimestampMixin:
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now()
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
    )


class SoftDeleteMixin:
    deleted_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True, index=True
    )


async def init_database() -> None:
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)


async def close_database() -> None:
    await engine.dispose()
