"""统一数据库连接模块 - 异步引擎与会话工厂"""

import logging
from collections.abc import AsyncGenerator

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

from app.config import settings

logger = logging.getLogger(__name__)

engine = create_async_engine(
    settings.database_url,
    echo=False,
    pool_pre_ping=True,
)

async_session_factory = async_sessionmaker(
    engine,
    class_=AsyncSession,
    expire_on_commit=False,
)


class Base(DeclarativeBase):
    """所有 ORM 模型的共享基类。"""

    pass


async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """FastAPI 依赖注入 - 提供异步数据库会话。"""
    async with async_session_factory() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise


async def init_db() -> None:
    """创建所有数据库表。"""
    from app.models import (  # noqa: F401 - 确保所有模型已导入
        AIAnalysisCache,
        AIModelConfig,
        AnalysisSettings,
        FuturesAnalysis,
        MarketData,
        ReviewDate,
        ReviewPlanV2,
        ScheduledTask,
        SectorHeat,
        Session,
        SymbolRanking,
        SymbolScoreV2,
        SymbolTimestamp,
        TradeImportBatch,
        TradeRecord,
        TradingPlan,
        TradingPlanV2,
        User,
        WatchedSymbol,
    )

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

    logger.info("数据库表创建完成")