"""
Checkpoint 工厂：PostgreSQL 优先，开发可降级内存。

本文件职责：按 CHECKPOINT_BACKEND / PG_DSN 创建 LangGraph checkpointer。
副作用：postgres 模式会 setup() 建 checkpoint 表（仅当前 PG_DATABASE）。
禁止：把 checkpoint 写入 Redis；禁止指向现网 ytd_runtime。
线程：Worker 进程启动时创建一次，供多槽共享（同一 thread_id 写冲突由 runtime 锁）。
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Any, Optional

from agent.config import Settings

logger = logging.getLogger(__name__)


@dataclass
class CheckpointBundle:
    """
    checkpointer 与可选连接资源。

    postgres 模式下保留 context/cm 以便进程退出时关闭连接池。
    """

    checkpointer: Any
    backend: str
    _closer: Optional[Any] = None

    def close(self) -> None:
        """释放 postgres 连接（若有）。无副作用于业务数据。"""
        if self._closer is None:
            return
        try:
            # from_conn_string 返回上下文管理器
            self._closer.__exit__(None, None, None)
        except Exception:  # noqa: BLE001 — 关闭失败只记日志
            logger.exception("关闭 Postgres checkpointer 失败")
        self._closer = None

def create_checkpointer(settings: Settings) -> CheckpointBundle:
    """
    创建 checkpointer。

    - CHECKPOINT_BACKEND=memory：仅本地骨架冒烟，进程重启丢书签。
    - postgres：需要可连的 resolve_pg_dsn()；失败时若 ytd_env=test 可降级 memory 并打警告。
    """
    backend = (settings.checkpoint_backend or "postgres").strip().lower()
    if backend == "memory":
        from langgraph.checkpoint.memory import MemorySaver

        logger.warning("Graph checkpoint 使用 MemorySaver（仅骨架/本机；正式必须 postgres）")
        return CheckpointBundle(checkpointer=MemorySaver(), backend="memory")

    dsn = settings.resolve_pg_dsn()
    if not dsn:
        if (settings.ytd_env or "").lower() == "prod":
            raise RuntimeError("正式环境必须配置 PG_DSN 或 PG_HOST/PG_DATABASE/PG_USER")
        from langgraph.checkpoint.memory import MemorySaver

        logger.warning(
            "未配置 PG_DSN，ytd_env=%s 降级 MemorySaver；正式切流前必须接 inquiry_robot_runtime",
            settings.ytd_env,
        )
        return CheckpointBundle(checkpointer=MemorySaver(), backend="memory")

    try:
        from langgraph.checkpoint.postgres import PostgresSaver

        # from_conn_string 返回上下文管理器；进入后得到 saver，并需保持打开
        cm = PostgresSaver.from_conn_string(dsn)
        saver = cm.__enter__()
        saver.setup()
        logger.info(
            "Graph Postgres checkpointer 就绪 database=%s",
            settings.pg_database or "(from DSN)",
        )
        return CheckpointBundle(checkpointer=saver, backend="postgres", _closer=cm)
    except Exception as exc:  # noqa: BLE001
        if (settings.ytd_env or "").lower() == "prod":
            raise
        from langgraph.checkpoint.memory import MemorySaver

        logger.warning(
            "Postgres checkpointer 初始化失败，降级 MemorySaver：%s",
            exc,
        )
        return CheckpointBundle(checkpointer=MemorySaver(), backend="memory")