"""
数据库连接模块
支持 PostgreSQL（生产）和 SQLite（本地开发）两种后端
"""

from __future__ import annotations

from collections.abc import AsyncGenerator
from pathlib import Path

from sqlalchemy import text
from sqlalchemy.ext.asyncio import (
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)
from sqlalchemy.pool import NullPool

from app.config import settings

# ──────────────────────────── 判断数据库类型 ────────────────────────────

DB_URL = settings.DATABASE_URL
IS_SQLITE = DB_URL.startswith("sqlite")

# ──────────────────────────── 异步引擎 ────────────────────────────

if IS_SQLITE:
    # SQLite 配置
    sqlite_path = DB_URL.replace("sqlite+aiosqlite:///", "")
    if not sqlite_path.startswith("/"):
        sqlite_path = str(Path(__file__).parent.parent / sqlite_path)

    engine = create_async_engine(
        f"sqlite+aiosqlite:///{sqlite_path}",
        echo=False,
    )
else:
    # PostgreSQL 配置
    engine = create_async_engine(
        DB_URL,
        echo=False,
        poolclass=NullPool,
        pool_pre_ping=True,
    )

# ──────────────────────────── 会话工厂 ────────────────────────────

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


# ──────────────────────────── 依赖注入 ────────────────────────────

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.base import Base

    async with engine.begin() as conn:
        if IS_SQLITE:
            # SQLite 启用 WAL 模式和外键约束
            await conn.execute(text("PRAGMA journal_mode=WAL"))
            await conn.execute(text("PRAGMA foreign_keys=ON"))
        # 创建所有表
        await conn.run_sync(Base.metadata.create_all)

    if not IS_SQLITE:
        # PostgreSQL 需要创建扩展
        async with engine.begin() as conn:
            try:
                await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
            except Exception:
                pass
            try:
                await conn.execute(text("CREATE EXTENSION IF NOT EXISTS zhparser"))
            except Exception:
                pass
            try:
                await conn.execute(text(
                    "CREATE TEXT SEARCH CONFIGURATION IF NOT EXISTS chinese_zh (PARSER = zhparser)"
                ))
            except Exception:
                pass


async def close_db() -> None:
    """关闭数据库引擎"""
    await engine.dispose()