"""Async checkpointer factory.

Provides an **async context manager** for long-running async servers that need
proper resource cleanup.

Supported backends: memory, sqlite, postgres.

Usage (e.g. FastAPI lifespan)::

    from deerflow.runtime.checkpointer.async_provider import make_checkpointer

    async with make_checkpointer() as checkpointer:
        app.state.checkpointer = checkpointer  # InMemorySaver if not configured

For sync usage see :mod:`deerflow.runtime.checkpointer.provider`.
"""

from __future__ import annotations

import asyncio
import contextlib
import logging
from collections.abc import AsyncIterator
import os
from pathlib import Path

from langgraph.types import Checkpointer

from deerflow.config.app_config import AppConfig, get_app_config
from deerflow.runtime.checkpointer.provider import (
    POSTGRES_CONN_REQUIRED,
    POSTGRES_INSTALL,
    SQLITE_INSTALL,
)
from deerflow.runtime.store._sqlite_utils import ensure_sqlite_parent_dir, resolve_sqlite_conn_str

logger = logging.getLogger(__name__)


def _csv_env(name: str) -> list[str]:
    raw = os.getenv(name, "")
    return [part.strip() for part in raw.split(",") if part.strip()]


def _dual_read_enabled(config) -> bool:
    if getattr(config, "dual_read_enabled", False):
        return True
    return os.getenv("DEER_FLOW_CHECKPOINT_DUAL_READ", "").strip().lower() in {"1", "true", "yes", "on"}


def _use_postgres_enabled(config) -> bool:
    raw = os.getenv("DEER_FLOW_USE_PG_CHECKPOINTER")
    if raw:
        return raw.strip().lower() not in {"0", "false", "no", "off"}
    return bool(getattr(config, "use_postgres", True))


def _dual_read_paths(config) -> list[str]:
    paths = list(getattr(config, "dual_read_sqlite_paths", []) or [])
    paths.extend(_csv_env("DEER_FLOW_CHECKPOINT_SQLITE_PATHS"))
    resolved: list[str] = []
    seen: set[str] = set()
    for path in paths:
        conn_str = resolve_sqlite_conn_str(path)
        if conn_str == ":memory:" or conn_str.startswith("file:"):
            logger.warning("Skipping unsupported checkpoint migration SQLite path: %s", path)
            continue
        if conn_str not in seen:
            resolved.append(conn_str)
            seen.add(conn_str)
    return resolved


def _sqlite_conn_str_for_postgres_fallback(config) -> str:
    configured = getattr(config, "sqlite_connection_string", None)
    if configured:
        return configured
    paths = _dual_read_paths(config)
    return paths[0] if paths else "store.db"


@contextlib.asynccontextmanager
async def _pg_pool_saver(conn_str: str) -> AsyncIterator[Checkpointer]:
    """Pool-backed ``AsyncPostgresSaver``, resilient to idle-closed connections.

    ``AsyncPostgresSaver.from_conn_string`` opens **one** long-lived connection.
    对远程 Postgres(本项目 47.88.25.99 这类),空闲几分钟后连接会被服务端 / NAT /
    连接池中间件断开,而 psycopg **不会自动重连** —— 下一次 checkpoint 操作就抛
    ``OperationalError: the connection is closed``,整个 run 直接失败(实测:圆桌第二步
    完成后闲置 ~6 分钟,用户再人工干预 → leader run 读 checkpoint 时撞死连接 → 前端
    "(无文本输出)")。

    用带健康检查的连接池根治:``check=AsyncConnectionPool.check_connection`` 在把连接交给
    调用方前先验活,坏连接直接丢弃重建;``max_idle`` 主动回收久闲连接,赶在远程 idle 超时
    之前。这样上层永远拿不到失效连接。``prepare_threshold=0`` 关预备语句,兼容 pgbouncer
    一类连接池(8432 非标准端口很可能就是)。
    """
    try:
        from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
    except ImportError as exc:
        raise ImportError(POSTGRES_INSTALL) from exc
    from psycopg_pool import AsyncConnectionPool

    try:
        max_size = int(os.getenv("DEER_FLOW_PG_POOL_MAX", "20"))
    except ValueError:
        max_size = 20

    pool = AsyncConnectionPool(
        conninfo=conn_str,
        min_size=1,
        max_size=max_size,
        open=False,
        kwargs={"autocommit": True, "prepare_threshold": 0},
        check=AsyncConnectionPool.check_connection,
        max_idle=120.0,
        max_lifetime=1800.0,
    )
    await pool.open(wait=True)
    try:
        saver = AsyncPostgresSaver(pool)
        await saver.setup()
        logger.info("Checkpointer: Postgres connection pool ready (max_size=%d, health-checked)", max_size)
        yield saver
    finally:
        await pool.close()


@contextlib.asynccontextmanager
async def _sqlite_checkpointer_from_conn_str(raw: str) -> AsyncIterator[Checkpointer]:
    try:
        from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
    except ImportError as exc:
        raise ImportError(SQLITE_INSTALL) from exc

    conn_str = resolve_sqlite_conn_str(raw)
    await asyncio.to_thread(ensure_sqlite_parent_dir, conn_str)
    async with AsyncSqliteSaver.from_conn_string(conn_str) as saver:
        await saver.setup()
        yield saver


# ---------------------------------------------------------------------------
# Async factory
# ---------------------------------------------------------------------------


@contextlib.asynccontextmanager
async def _async_checkpointer(config) -> AsyncIterator[Checkpointer]:
    """Async context manager that constructs and tears down a checkpointer."""
    if config.type == "memory":
        from langgraph.checkpoint.memory import InMemorySaver

        yield InMemorySaver()
        return

    if config.type == "sqlite":
        async with _sqlite_checkpointer_from_conn_str(config.connection_string or "store.db") as saver:
            yield saver
        return

    if config.type == "postgres":
        if not _use_postgres_enabled(config):
            conn_str = _sqlite_conn_str_for_postgres_fallback(config)
            logger.info("Checkpointer: use_postgres=false; using SQLite fallback (%s)", conn_str)
            async with _sqlite_checkpointer_from_conn_str(conn_str) as saver:
                yield saver
            return

        if not config.connection_string:
            raise ValueError(POSTGRES_CONN_REQUIRED)

        # 连接池 + 健康检查,根治远程 Postgres 单连接空闲被关导致的 run 崩溃(见 _pg_pool_saver)。
        async with _pg_pool_saver(config.connection_string) as saver:
            if not _dual_read_enabled(config):
                yield saver
                return

            from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver

            from deerflow.runtime.checkpointer.checkpoint_migration import (
                DualReadCheckpointer,
                SQLiteCheckpointSource,
            )

            sources: list[SQLiteCheckpointSource] = []
            for sqlite_path in _dual_read_paths(config):
                if not Path(sqlite_path).exists():
                    logger.info("Checkpoint dual-read source does not exist, skipping: %s", sqlite_path)
                    continue
                ctx = AsyncSqliteSaver.from_conn_string(sqlite_path)
                source_saver = await ctx.__aenter__()
                await source_saver.setup()
                sources.append(
                    SQLiteCheckpointSource(
                        path=sqlite_path,
                        saver=source_saver,
                        context=ctx,
                        progress_path=Path(sqlite_path).with_suffix(Path(sqlite_path).suffix + ".pg-migration.json"),
                    )
                )

            if not sources:
                logger.info("Checkpoint dual-read enabled but no SQLite sources were opened; using Postgres only")
                yield saver
                return

            proxy = DualReadCheckpointer(
                saver,
                sources,
            )
            try:
                yield proxy
            finally:
                await proxy.aclose_sources()
        return

    raise ValueError(f"Unknown checkpointer type: {config.type!r}")


# ---------------------------------------------------------------------------
# Public async context manager
# ---------------------------------------------------------------------------


@contextlib.asynccontextmanager
async def _async_checkpointer_from_database(db_config) -> AsyncIterator[Checkpointer]:
    """Async context manager that constructs a checkpointer from unified DatabaseConfig."""
    if db_config.backend == "memory":
        from langgraph.checkpoint.memory import InMemorySaver

        yield InMemorySaver()
        return

    if db_config.backend == "sqlite":
        try:
            from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
        except ImportError as exc:
            raise ImportError(SQLITE_INSTALL) from exc

        conn_str = db_config.checkpointer_sqlite_path
        ensure_sqlite_parent_dir(conn_str)
        async with AsyncSqliteSaver.from_conn_string(conn_str) as saver:
            await saver.setup()
            yield saver
        return

    if db_config.backend == "postgres":
        if not db_config.postgres_url:
            raise ValueError("database.postgres_url is required for the postgres backend")

        # 连接池 + 健康检查(见 _pg_pool_saver),避免单连接空闲被关。
        async with _pg_pool_saver(db_config.postgres_url) as saver:
            yield saver
        return

    raise ValueError(f"Unknown database backend: {db_config.backend!r}")


@contextlib.asynccontextmanager
async def make_checkpointer(app_config: AppConfig | None = None) -> AsyncIterator[Checkpointer]:
    """Async context manager that yields a checkpointer for the caller's lifetime.
    Resources are opened on enter and closed on exit -- no global state::

        async with make_checkpointer(app_config) as checkpointer:
            app.state.checkpointer = checkpointer

    Yields an ``InMemorySaver`` when no checkpointer is configured in *config.yaml*.

    Priority:
    1. Legacy ``checkpointer:`` config section (backward compatible)
    2. Unified ``database:`` config section
    3. Default InMemorySaver
    """

    if app_config is None:
        app_config = get_app_config()

    # Legacy: standalone checkpointer config takes precedence
    if app_config.checkpointer is not None:
        async with _async_checkpointer(app_config.checkpointer) as saver:
            yield saver
            return

    # Unified database config
    db_config = getattr(app_config, "database", None)
    if db_config is not None and db_config.backend != "memory":
        async with _async_checkpointer_from_database(db_config) as saver:
            yield saver
            return

    # Default: in-memory
    from langgraph.checkpoint.memory import InMemorySaver

    yield InMemorySaver()