"""Simple PostgreSQL connectivity smoke-checker."""
import os
import sys
import psycopg2
from psycopg2 import OperationalError


DEFAULT_DSN = os.environ.get(
    "PG_DSN", "postgresql://local-Python:Neo-local-1991125@100.64.0.4:5432/LLZQ-test"
)
DEFAULT_TIMEOUT = max(1, min(int(os.environ.get("PG_CONNECT_TIMEOUT", 10)), 20))


def check_postgres_connection(dsn: str, timeout: int = DEFAULT_TIMEOUT) -> bool:
    """Return True if connection succeeds; print diagnostics otherwise."""
    try:
        conn = psycopg2.connect(dsn, connect_timeout=timeout)
        with conn:
            with conn.cursor() as cur:
                cur.execute("SELECT 1;")
                _ = cur.fetchone()
        print(f"PostgreSQL 连接成功 (timeout={timeout}s)")
        return True
    except OperationalError as exc:
        print("PostgreSQL 连接失败（OperationalError）：", exc)
    except Exception as exc:  # pragma: no cover - defensive
        print("PostgreSQL 连接失败（其他异常）：", exc)

    return False


if __name__ == "__main__":
    dsn = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_DSN
    if not dsn:
        print("缺少 DSN，请传入参数或设置 PG_DSN 环境变量。")
        sys.exit(2)

    ok = check_postgres_connection(dsn)
    if not ok:
        sys.exit(1)