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 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 才建议配置
if _is_celery_process() and not settings.DATABASE_URL.startswith("sqlite"):
    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()