from __future__ import annotations

import threading

import psycopg
from langgraph.checkpoint.postgres import PostgresSaver
from psycopg.rows import dict_row

from app.modules.shared.db import database_url

_CHECKPOINTER: PostgresSaver | None = None
_LOCK = threading.Lock()


def _conn_string() -> str:
    url = database_url()
    if url.startswith("postgresql+psycopg"):
        return url.replace("postgresql+psycopg", "postgresql", 1)
    return url


def get_checkpointer() -> PostgresSaver:
    global _CHECKPOINTER
    with _LOCK:
        if _CHECKPOINTER is None:
            conn = psycopg.connect(_conn_string(), autocommit=True, row_factory=dict_row)
            cp = PostgresSaver(conn)
            cp.setup()
            _CHECKPOINTER = cp
        return _CHECKPOINTER