"""数据库连接管理"""
import asyncio
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.orm import DeclarativeBase
from app.core.config import settings
from app.core.logger import logger


class Base(DeclarativeBase):
    pass


engine = create_async_engine(
    settings.DATABASE_URL,
    echo=False,
    pool_pre_ping=True,
    pool_recycle=3600,
    pool_size=10,
    max_overflow=20,
)

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


async def get_db():
    """获取数据库会话"""
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()


async def wait_for_db(max_retries: int = 30, interval: int = 2):
    """等待 MySQL 就绪，最多重试 max_retries 次"""
    for attempt in range(1, max_retries + 1):
        try:
            async with engine.begin() as conn:
                await conn.execute(__import__("sqlalchemy").text("SELECT 1"))
            logger.info(f"✅ 数据库连接成功（第 {attempt} 次尝试）")
            return
        except Exception as e:
            if attempt == max_retries:
                logger.error(f"数据库连接失败，已重试 {max_retries} 次: {e}")
                raise
            logger.warning(f"数据库未就绪，{interval}s 后重试（{attempt}/{max_retries}）: {e}")
            await asyncio.sleep(interval)


async def init_db():
    """初始化数据库 - 等待 MySQL 就绪并注册所有模型"""
    try:
        # 等待 MySQL 容器真正就绪
        await wait_for_db(max_retries=30, interval=2)

        # 导入所有模型类，确保 SQLAlchemy ORM 元数据注册
        from app.models import (
            VirtualUser, UserPersonality, InteractionRecord,
            TokenStat, AIModelConfig, SystemConfig, LoginLog
        )
        logger.info("✅ 数据库模型注册成功")
        logger.info("✅ 数据库初始化完成")

    except Exception as e:
        logger.error(f"数据库初始化失败: {e}")
        raise