from __future__ import annotations

import os
from collections.abc import Iterator
from contextlib import contextmanager


@contextmanager
def open_story_checkpointer() -> Iterator[object]:
    database_url = os.environ.get("DATABASE_URL", "")
    if database_url.startswith(("postgres://", "postgresql://")):
        try:
            from langgraph.checkpoint.postgres import PostgresSaver
        except ImportError as exc:
            raise RuntimeError(
                "Spark story workflows require langgraph-checkpoint-postgres; install project dependencies"
            ) from exc
        with PostgresSaver.from_conn_string(database_url) as saver:
            saver.setup()
            yield saver
        return
    from langgraph.checkpoint.memory import MemorySaver

    yield MemorySaver()