"""持久化装配：按配置选择 MongoDB 或进程内存，统一返回会话存储、检查点和关闭函数"""

from collections.abc import Awaitable, Callable
from dataclasses import dataclass

from langgraph.checkpoint.base import BaseCheckpointSaver
from loguru import logger

from app.config import settings
from app.services.conversation import ConversationStore, InMemoryConversationStore


@dataclass
class Persistence:
    store: ConversationStore
    # 为空时由图使用默认的内存检查点
    checkpointer: BaseCheckpointSaver | None
    close: Callable[[], Awaitable[None]]


async def _close_nothing() -> None:
    return None


async def open_persistence() -> Persistence:
    if not settings.mongodb_uri:
        logger.warning(
            "未配置 MONGODB_URI，会话记录和检查点仅保存在进程内存，重启后丢失"
        )
        return Persistence(InMemoryConversationStore(), None, _close_nothing)

    # 依赖较重且只在启用 MongoDB 时使用，放到函数内导入
    from langgraph.checkpoint.mongodb import MongoDBSaver
    from pymongo import AsyncMongoClient, MongoClient

    from app.processor.graph import build_checkpoint_serde
    from app.services.conversation_mongo import MongoConversationStore

    # 会话记录使用异步客户端；官方检查点基于同步客户端，异步方法由其放入线程池执行
    async_client = AsyncMongoClient(settings.mongodb_uri, serverSelectionTimeoutMS=5000)
    sync_client = MongoClient(settings.mongodb_uri, serverSelectionTimeoutMS=5000)

    async def close() -> None:
        try:
            await async_client.close()
        finally:
            sync_client.close()

    try:
        # 启动时探测连接，配置错误尽早暴露；异常信息不包含连接串，避免泄露账号密码
        try:
            await async_client.admin.command("ping")
        except Exception as exc:
            raise RuntimeError(f"无法连接 MongoDB：{type(exc).__name__}") from None

        store = MongoConversationStore(async_client[settings.mongodb_db])
        await store.ensure_indexes()
        # 与内存检查点共用同一份类型白名单，保证读回后仍是业务模型对象
        checkpointer = MongoDBSaver(
            sync_client,
            settings.mongodb_db,
            ttl=settings.mongodb_checkpoint_ttl or None,
            serde=build_checkpoint_serde(),
        )
    except BaseException:
        await close()
        raise

    logger.info("已启用 MongoDB 持久化，库名 {}", settings.mongodb_db)
    return Persistence(store, checkpointer, close)