#!/usr/bin/env python3
"""
03_logical_replication.py
=========================
演示 PG 逻辑复制：在主库建 PUBLICATION，在从库建 SUBSCRIPTION，
然后在主库写数据，验证从库能同步看到。

前置条件：
    - 主库（端口 5432）和从库（端口 5433，独立实例非物理 standby）都在跑
    - 两边 wal_level >= logical
    - 两边都有 learn_pg 数据库
    - repl_user 账号已创建
    - 两边都有 ch15_logical_demo 表（脚本会自动建）

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

用法：
    python3 03_logical_replication.py setup    # 建表 + 发布 + 订阅
    python3 03_logical_replication.py write    # 主库写入测试
    python3 03_logical_replication.py verify   # 验证从库同步
    python3 03_logical_replication.py status   # 查 publication / subscription
    python3 03_logical_replication.py teardown # 清理
"""
import os
import sys
import time

import psycopg

MASTER_DSN = os.environ.get(
    "MASTER_DSN",
    "host=127.0.0.1 port=5432 dbname=learn_pg user=postgres password=postgres",
)
REPLICA_DSN = os.environ.get(
    "REPLICA_DSN",
    "host=127.0.0.1 port=5433 dbname=learn_pg user=postgres password=postgres",
)
# subscription 端连接主库用的 DSN（注意密码改成 repl_user 的）
SUB_CONN = os.environ.get(
    "SUB_CONN",
    "host=127.0.0.1 port=5432 dbname=learn_pg user=repl_user password=r3pl_pwd",
)


def exec_sql(dsn: str, sql: str, fetch: bool = False):
    with psycopg.connect(dsn, autocommit=True) as conn, conn.cursor() as cur:
        cur.execute(sql)
        if fetch and cur.description:
            return cur.fetchall()
    return None


def setup() -> None:
    print("[1/4] 主库建表 + 发布")
    exec_sql(MASTER_DSN, """
        CREATE TABLE IF NOT EXISTS ch15_logical_demo (
            id     BIGSERIAL PRIMARY KEY,
            payload TEXT NOT NULL,
            ts     TIMESTAMPTZ DEFAULT now()
        );
        DROP PUBLICATION IF EXISTS ch15_pub_logical_demo;
        CREATE PUBLICATION ch15_pub_logical_demo FOR TABLE ch15_logical_demo;
    """)

    print("[2/4] 从库建表（DDL 不会自动复制，必须手工建）")
    exec_sql(REPLICA_DSN, """
        CREATE TABLE IF NOT EXISTS ch15_logical_demo (
            id     BIGSERIAL PRIMARY KEY,
            payload TEXT NOT NULL,
            ts     TIMESTAMPTZ DEFAULT now()
        );
        TRUNCATE ch15_logical_demo;
    """)

    print("[3/4] 从库建订阅")
    exec_sql(REPLICA_DSN, "DROP SUBSCRIPTION IF EXISTS ch15_sub_logical_demo;")
    exec_sql(REPLICA_DSN, f"""
        CREATE SUBSCRIPTION ch15_sub_logical_demo
        CONNECTION '{SUB_CONN}'
        PUBLICATION ch15_pub_logical_demo;
    """)

    print("[4/4] 等 5 秒让初始 sync 完成...")
    time.sleep(5)
    print("✅ 设置完成。")


def write() -> None:
    print("--- 主库插入 10 行数据 ---")
    rows = exec_sql(
        MASTER_DSN,
        """
        INSERT INTO ch15_logical_demo(payload)
        SELECT 'logical-' || g FROM generate_series(1,10) g
        RETURNING id, payload;
        """,
        fetch=True,
    )
    for r in rows or []:
        print(f"  inserted: {r}")


def verify() -> None:
    print("--- 主库当前数据量 ---")
    m = exec_sql(MASTER_DSN, "SELECT count(*), max(id) FROM ch15_logical_demo;", fetch=True)
    print(f"  master: count={m[0][0]} max_id={m[0][1]}")

    print("--- 等 2 秒让逻辑复制追上 ---")
    time.sleep(2)

    print("--- 从库当前数据量 ---")
    r = exec_sql(REPLICA_DSN, "SELECT count(*), max(id) FROM ch15_logical_demo;", fetch=True)
    print(f"  replica: count={r[0][0]} max_id={r[0][1]}")

    if m[0][0] == r[0][0]:
        print("✅ 主从一致")
    else:
        print("⚠️  主从不一致，请检查 pg_stat_subscription / pg_stat_replication")


def status() -> None:
    print("=== 主库 pg_publication ===")
    for row in exec_sql(MASTER_DSN, "SELECT pubname, puballtables FROM pg_publication;", fetch=True) or []:
        print(f"  {row}")

    print("\n=== 主库 pg_replication_slots（逻辑槽） ===")
    for row in exec_sql(
        MASTER_DSN,
        "SELECT slot_name, slot_type, plugin, active FROM pg_replication_slots WHERE slot_type='logical';",
        fetch=True,
    ) or []:
        print(f"  {row}")

    print("\n=== 主库 pg_stat_replication ===")
    for row in exec_sql(
        MASTER_DSN,
        "SELECT application_name, client_addr, state, sync_state FROM pg_stat_replication;",
        fetch=True,
    ) or []:
        print(f"  {row}")

    print("\n=== 从库 pg_subscription ===")
    for row in exec_sql(
        REPLICA_DSN,
        "SELECT subname, subenabled, subslotname FROM pg_subscription;",
        fetch=True,
    ) or []:
        print(f"  {row}")

    print("\n=== 从库 pg_stat_subscription ===")
    for row in exec_sql(
        REPLICA_DSN,
        "SELECT subname, pid, received_lsn, latest_end_lsn FROM pg_stat_subscription;",
        fetch=True,
    ) or []:
        print(f"  {row}")


def teardown() -> None:
    print("[1/3] 从库删订阅")
    try:
        exec_sql(REPLICA_DSN, "ALTER SUBSCRIPTION ch15_sub_logical_demo DISABLE;")
        exec_sql(REPLICA_DSN, "ALTER SUBSCRIPTION ch15_sub_logical_demo SET (slot_name = NONE);")
        exec_sql(REPLICA_DSN, "DROP SUBSCRIPTION ch15_sub_logical_demo;")
    except psycopg.Error as e:
        print(f"  warn: {e}")

    print("[2/3] 主库删发布 + 槽")
    try:
        exec_sql(MASTER_DSN, "DROP PUBLICATION IF EXISTS ch15_pub_logical_demo;")
        exec_sql(MASTER_DSN, "SELECT pg_drop_replication_slot('ch15_sub_logical_demo') WHERE EXISTS (SELECT 1 FROM pg_replication_slots WHERE slot_name='ch15_sub_logical_demo');")
    except psycopg.Error as e:
        print(f"  warn: {e}")

    print("[3/3] 删表")
    exec_sql(MASTER_DSN, "DROP TABLE IF EXISTS ch15_logical_demo;")
    exec_sql(REPLICA_DSN, "DROP TABLE IF EXISTS ch15_logical_demo;")
    print("✅ 已清理。")


COMMANDS = {
    "setup": setup,
    "write": write,
    "verify": verify,
    "status": status,
    "teardown": teardown,
}

if __name__ == "__main__":
    cmd = sys.argv[1] if len(sys.argv) > 1 else ""
    if cmd not in COMMANDS:
        print(f"Usage: {sys.argv[0]} {{{'|'.join(COMMANDS.keys())}}}", file=sys.stderr)
        sys.exit(1)
    COMMANDS[cmd]()