"""Tortoise ORM 初始化"""
import logging
from tortoise import Tortoise, connections
import config

logger = logging.getLogger("spam_guard")

MODELS = [
    "db.models",
]

# aerich 配置
TORTOISE_ORM = {
    "connections": {
        "default": config.DATABASE_URL,
    },
    "apps": {
        "models": {
            "models": MODELS,
            "default_connection": "default",
        },
    },
}


async def ensure_runtime_schema():
    """轻量补齐运行期新增列；只做 ADD COLUMN，不破坏已有数据。"""
    conn = connections.get("default")
    try:
        await conn.execute_script("""
        ALTER TABLE messages ADD COLUMN IF NOT EXISTS manually_deleted BOOLEAN NOT NULL DEFAULT FALSE;
        ALTER TABLE messages ADD COLUMN IF NOT EXISTS deleted_at TIMESTAMPTZ NULL;
        ALTER TABLE messages ADD COLUMN IF NOT EXISTS delete_reason VARCHAR(255) NULL;
        ALTER TABLE user_bindings ADD COLUMN IF NOT EXISTS auth_email VARCHAR(255) NULL;
        ALTER TABLE user_bindings ADD COLUMN IF NOT EXISTS time_expire VARCHAR(100) NULL;
        """)
    except Exception as e:
        logger.warning(f"⚠️ 自动补齐数据库字段失败: {type(e).__name__}: {e}")


async def init_db():
    await Tortoise.init(
        db_url=config.DATABASE_URL,
        modules={"models": MODELS},
    )
    # aerich 用 generate_schemas；正式环境用 aerich migrate
    await Tortoise.generate_schemas()
    await ensure_runtime_schema()
    logger.info("✅ 数据库连接成功")


async def close_db():
    await Tortoise.close_connections()
    logger.info("🛑 数据库连接已关闭")