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

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

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]]
    # 启用 MongoDB 时的数据库，供物料等其他集合复用同一连接
    db: Any = 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, async_client[settings.mongodb_db])