"""
04_long_transaction_pain.py —— 长事务阻止 VACUUM 回收死元组

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

实验设计：
    1. 主线程：开一个 REPEATABLE READ 长事务并持有快照
    2. 后台线程：反复 UPDATE 同一批行，制造死元组
    3. 主线程发起 VACUUM (VERBOSE)
       -> 输出会有 "0 are dead but not yet removable" 字样
    4. 主线程提交长事务后再 VACUUM
       -> 死元组终于被回收
"""
from __future__ import annotations

import threading
import time

import psycopg
from psycopg import IsolationLevel

DSN = "host=127.0.0.1 port=5432 dbname=learn_pg user=postgres"
TABLE = "ch8_long_tx_demo"


def setup() -> None:
    with psycopg.connect(DSN, autocommit=True) as conn:
        conn.execute(f"TRUNCATE {TABLE}")
        conn.execute(
            f"INSERT INTO {TABLE} SELECT g, g FROM generate_series(1, 1000) g"
        )
        conn.execute(f"VACUUM FULL {TABLE}")


def stats(label: str) -> None:
    with psycopg.connect(DSN, autocommit=True) as conn:
        conn.execute(f"ANALYZE {TABLE}")
        cur = conn.execute("""
            SELECT n_live_tup, n_dead_tup,
                   pg_size_pretty(pg_relation_size(%s))
            FROM pg_stat_user_tables WHERE relname='ch8_long_tx_demo'
        """, (TABLE,))
        live, dead, size = cur.fetchone()
    print(f"  [{label:<32}] live={live:>5} dead={dead:>5} size={size}")


def churn() -> None:
    """后台线程：制造死元组"""
    with psycopg.connect(DSN, autocommit=True) as conn:
        for _ in range(20):
            conn.execute(f"UPDATE {TABLE} SET v = v + 1 WHERE id BETWEEN 1 AND 500")


def vacuum_verbose(tag: str) -> None:
    """运行 VACUUM (VERBOSE) 并捕获 NOTICE 输出。"""
    print(f"\n--- VACUUM {tag} ---")
    captured: list[str] = []
    with psycopg.connect(DSN, autocommit=True) as conn:
        # psycopg v3：通过 notice handler 收集 NOTICE/INFO
        conn.add_notice_handler(
            lambda diag: captured.append(diag.message_primary or "")
        )
        with conn.cursor() as cur:
            cur.execute(f"VACUUM (VERBOSE) {TABLE}")
    for note in captured:
        if any(kw in note for kw in
               ("removable", "tuples:", "pages:", "scan", "removed",
                "remain", "frozen", "index")):
            print(f"    {note.strip()}")


def main() -> None:
    setup()
    print("=== 第一阶段：没有长事务，VACUUM 正常工作 ===")
    stats("初始")
    churn()
    stats("UPDATE 后")
    vacuum_verbose("（无长事务干扰）")
    stats("VACUUM 后")

    print("\n=== 第二阶段：开一个长事务持有 snapshot ===")
    long_conn = psycopg.connect(DSN)
    long_conn.isolation_level = IsolationLevel.REPEATABLE_READ
    long_conn.execute(f"SELECT count(*) FROM {TABLE}")  # 触发 snapshot
    print("  长事务已开启并持有 snapshot")

    setup()
    churn()
    stats("制造死元组（长事务在跑）")

    vacuum_verbose("（长事务正在跑，应大量 not yet removable）")
    stats("VACUUM 后（死元组没回收）")

    print("\n=== 第三阶段：长事务 COMMIT 后 VACUUM 才有效 ===")
    long_conn.commit()
    long_conn.close()
    print("  长事务已提交")

    vacuum_verbose("（长事务已结束，应能完整回收）")
    stats("VACUUM 后（终于回收）")

    print("\n💡 结论：长事务（特别是 idle in transaction）会让 PG 表无限膨胀！")
    print("    生产必须设置 idle_in_transaction_session_timeout 并监控告警。")


if __name__ == "__main__":
    main()