#!/usr/bin/env python3
"""
01_show_replication_status.py
=============================
从主库查询当前所有从库的复制状态、复制槽状态、复制延迟。

依赖：psycopg[binary] >= 3.1
    pip install "psycopg[binary]"

连接：默认 host=127.0.0.1 port=5432 dbname=learn_pg user=postgres
"""
import os
import sys
import psycopg

CONN_INFO = os.environ.get(
    "PG_DSN",
    "host=127.0.0.1 port=5432 dbname=learn_pg user=postgres password=postgres",
)


def hr(title: str) -> None:
    print("\n" + "=" * 70)
    print(f"  {title}")
    print("=" * 70)


def query_and_print(cur, sql: str) -> None:
    cur.execute(sql)
    cols = [d.name for d in cur.description]
    rows = cur.fetchall()
    if not rows:
        print("  (无记录)")
        return
    widths = [max(len(c), max((len(str(r[i])) for r in rows), default=0)) for i, c in enumerate(cols)]
    fmt = "  " + "  ".join(f"{{:<{w}}}" for w in widths)
    print(fmt.format(*cols))
    print("  " + "-" * (sum(widths) + 2 * (len(cols) - 1)))
    for r in rows:
        print(fmt.format(*[str(v) for v in r]))


def main() -> int:
    try:
        conn = psycopg.connect(CONN_INFO)
    except psycopg.Error as e:
        print(f"[FATAL] 无法连接主库: {e}", file=sys.stderr)
        return 1

    with conn, conn.cursor() as cur:
        hr("1. 当前 PG 角色")
        cur.execute("SELECT pg_is_in_recovery() AS in_recovery, current_setting('wal_level') AS wal_level;")
        in_recovery, wal_level = cur.fetchone()
        print(f"  pg_is_in_recovery = {in_recovery}   {'(从库)' if in_recovery else '(主库)'}")
        print(f"  wal_level         = {wal_level}")

        hr("2. pg_stat_replication（连上来的从库）")
        query_and_print(cur, """
            SELECT
                application_name,
                client_addr,
                state,
                sync_state,
                pg_size_pretty(pg_wal_lsn_diff(pg_current_wal_lsn(), sent_lsn))    AS pending_send,
                pg_size_pretty(pg_wal_lsn_diff(sent_lsn, write_lsn))               AS pending_write,
                pg_size_pretty(pg_wal_lsn_diff(write_lsn, flush_lsn))              AS pending_flush,
                pg_size_pretty(pg_wal_lsn_diff(flush_lsn, replay_lsn))             AS pending_replay,
                pg_size_pretty(pg_wal_lsn_diff(pg_current_wal_lsn(), replay_lsn))  AS total_lag,
                COALESCE(write_lag::text,  '-') AS write_lag,
                COALESCE(flush_lag::text,  '-') AS flush_lag,
                COALESCE(replay_lag::text, '-') AS replay_lag
            FROM pg_stat_replication
            ORDER BY application_name;
        """)

        hr("3. pg_replication_slots（所有复制槽）")
        query_and_print(cur, """
            SELECT
                slot_name,
                slot_type,
                database,
                active,
                COALESCE(active_pid::text, '-') AS active_pid,
                pg_size_pretty(pg_wal_lsn_diff(pg_current_wal_lsn(), restart_lsn)) AS retained_wal
            FROM pg_replication_slots
            ORDER BY slot_name;
        """)

        hr("4. WAL 写入位置（主库视角）")
        cur.execute("""
            SELECT
                pg_current_wal_lsn() AS current_lsn,
                pg_walfile_name(pg_current_wal_lsn()) AS current_walfile;
        """)
        lsn, walfile = cur.fetchone()
        print(f"  current_lsn      = {lsn}")
        print(f"  current_walfile  = {walfile}")

        hr("5. PUBLICATIONS")
        query_and_print(cur, """
            SELECT pubname, puballtables, pubinsert, pubupdate, pubdelete, pubtruncate
            FROM pg_publication;
        """)

    print("\n[DONE] 查询完成。如要持续监控，可加上 watch 1：")
    print("  watch -n 1 \"python3 01_show_replication_status.py | head -50\"")
    return 0


if __name__ == "__main__":
    sys.exit(main())