"""langgraph checkpoint setup (multi-chapter chain workflow)

Revision ID: d3e4f5a6b7c8
Revises: c2b3d4e5f6a7
Create Date: 2026-06-23 10:00:00.000000

多章工作流链（docs/design/chain-workflow.md §4/§7）唯一新增 DDL：langgraph
检查点表（`checkpoints` / `checkpoint_blobs` / `checkpoint_writes` /
`checkpoint_migrations`）。这些表**不在** ``Base.metadata``（由 langgraph 自己
管理 schema），故迁移直接调用 langgraph 的一次性 ``setup()`` 建表，而非 ORM
``op.create_table``。env.py 的 ``include_object`` 过滤这些表，避免 ``alembic
check`` 误报漂移（要 drop 这些"非 metadata"表）。

实现说明（CLAUDE.md「LangGraph」+ 设计 §11 回退）：
- DDL 只在 **migrations/CI** 跑（绝不在 app-runtime）——本迁移即该入口。
- ``upgrade()`` 跑在 alembic 的**同步**迁移上下文里（``connection.run_sync``
  下已有 event loop 在跑 async 迁移），故**不能**调 async ``setup_checkpointer``
  /``asyncio.run``（会嵌套 loop 报错）。改用 langgraph 同步 ``PostgresSaver``
  （psycopg），用一条从 ``database_url`` 派生的 psycopg conn string 自建连接，
  与 alembic 的 asyncpg 连接解耦。``setup()`` 幂等（langgraph 内部 migration 版本表）。
- **关键**：langgraph ``setup()`` 跑 ``CREATE INDEX CONCURRENTLY``，它不能在
  事务块内执行，且会等 alembic 那条仍开着事务（持 ``alembic_version`` 锁）的连接
  → 死锁（§11 风险落地）。故用 ``op.get_context().autocommit_block()`` 把 alembic
  的事务先提交、退出事务块，再让 langgraph 在独立 autocommit 连接里建表/建索引。

jobs 表：C2 判定零迁移复用（用 ``result.awaiting_chapter`` 表达待裁决，
``status`` 复用现有枚举），故本迁移不动 jobs。
"""

from __future__ import annotations

from collections.abc import Sequence

revision: str = "d3e4f5a6b7c8"
down_revision: str | None = "c2b3d4e5f6a7"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None

# langgraph 检查点表（langgraph.checkpoint.postgres.base 的 MIGRATIONS 建的全部表）
_CHECKPOINT_TABLES: tuple[str, ...] = (
    "checkpoint_writes",
    "checkpoint_blobs",
    "checkpoints",
    "checkpoint_migrations",
)


def _psycopg_conn_string() -> str:
    """从 settings.database_url 派生 psycopg（同步）conn string。

    settings 用 asyncpg URL（``postgresql+asyncpg://…``）；langgraph 的同步
    ``PostgresSaver`` 走 psycopg，需去掉 SQLAlchemy 的 ``+asyncpg`` 驱动后缀。
    """
    from ww_config import get_settings

    return get_settings().database_url.replace("+asyncpg", "")


def upgrade() -> None:
    from alembic import op
    from langgraph.checkpoint.postgres import PostgresSaver

    # 退出 alembic 事务块（提交 alembic_version 更新，释放其锁），langgraph 的
    # CREATE INDEX CONCURRENTLY 才能在独立连接里跑而不死锁。
    with op.get_context().autocommit_block():
        with PostgresSaver.from_conn_string(_psycopg_conn_string()) as saver:
            saver.setup()


def downgrade() -> None:
    import psycopg

    with psycopg.connect(_psycopg_conn_string()) as conn:
        conn.execute("DROP TABLE IF EXISTS " + ", ".join(_CHECKPOINT_TABLES) + " CASCADE")
        conn.commit()