"""
04_wal_size_estimate.py
-----------------------
精确测量「批量 INSERT / UPDATE / DELETE / TRUNCATE」分别产生多少 WAL：
  · 用 pg_current_wal_lsn() 在操作前后取 LSN
  · pg_wal_lsn_diff 求字节差
  · 进一步对比开 / 关 wal_compression 时 FPI 的变化
"""

import psycopg

CONN_STR = "host=127.0.0.1 port=5432 dbname=learn_pg user=postgres"
N_ROWS = 10000


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


def lsn(cur: psycopg.Cursor) -> str:
    cur.execute("SELECT pg_current_wal_lsn();")
    return cur.fetchone()[0]


def diff(cur: psycopg.Cursor, a: str, b: str) -> int:
    cur.execute("SELECT pg_wal_lsn_diff(%s, %s);", (a, b))
    return int(cur.fetchone()[0])


def measure(cur: psycopg.Cursor, label: str, sql: str) -> int:
    cur.execute("CHECKPOINT;")     # 隔离掉之前的脏页
    before = lsn(cur)
    cur.execute(sql)
    after = lsn(cur)
    bytes_ = diff(cur, after, before)
    print(f"  {label:<40} = {bytes_:>10} B ≈ {bytes_/1024:>7.1f} KB")
    return bytes_


def main() -> None:
    with psycopg.connect(CONN_STR, autocommit=True) as conn, conn.cursor() as cur:
        cur.execute("DROP TABLE IF EXISTS ch10_ws_demo;")
        cur.execute(
            "CREATE TABLE ch10_ws_demo (id INT PRIMARY KEY, name TEXT, val INT);"
        )

        banner(f"1. 各种 DML 产生的 WAL 量（每批 {N_ROWS} 行）")
        measure(
            cur, "INSERT (批量)",
            f"INSERT INTO ch10_ws_demo (id, name, val) "
            f"SELECT g, 'name-' || g, g FROM generate_series(1, {N_ROWS}) g;"
        )
        measure(
            cur, "UPDATE 全表（非索引列）",
            "UPDATE ch10_ws_demo SET val = val + 1;"
        )
        measure(
            cur, "UPDATE 全表（被索引列：会触发索引页修改）",
            "UPDATE ch10_ws_demo SET id = id + 100000;"
        )
        # 还原便于后续对比
        cur.execute("UPDATE ch10_ws_demo SET id = id - 100000;")
        measure(
            cur, "DELETE 全表",
            "DELETE FROM ch10_ws_demo;"
        )
        measure(
            cur, "INSERT 重建",
            f"INSERT INTO ch10_ws_demo (id, name, val) "
            f"SELECT g, 'n-' || g, g FROM generate_series(1, {N_ROWS}) g;"
        )
        measure(
            cur, "TRUNCATE（不写元组日志，只 metadata）",
            "TRUNCATE ch10_ws_demo;"
        )
        measure(
            cur, "DROP + CREATE",
            "DROP TABLE ch10_ws_demo;"
            f"CREATE TABLE ch10_ws_demo (id INT PRIMARY KEY, name TEXT, val INT);"
            f"INSERT INTO ch10_ws_demo SELECT g, 'n-' || g, g "
            f"FROM generate_series(1, {N_ROWS}) g;"
        )

        # ---------- 2. wal_compression 对比 ----------
        banner("2. wal_compression 对 FPI 的影响")
        cur.execute("SHOW wal_compression;")
        original = cur.fetchone()[0]
        print(f"  当前 wal_compression = {original}")

        try:
            for setting in ("off", "on"):
                cur.execute(f"SET wal_compression = '{setting}';")
                # 重新填表后立刻 checkpoint，确保接下来的 UPDATE 全是 FPI 写
                cur.execute("CHECKPOINT;")
                size = measure(
                    cur,
                    f"UPDATE 全表（wal_compression = {setting:>3}）",
                    "UPDATE ch10_ws_demo SET val = val + 1;"
                )
            cur.execute(f"SET wal_compression = '{original}';")
        except psycopg.Error as e:
            print(f"  跳过：{e}")

        banner("3. 解读")
        print(
            """
  · INSERT 比 UPDATE 便宜：UPDATE 要同时记录老元组的 t_xmax 修改 + 新元组 + 索引项
  · UPDATE 索引列比改非索引列贵：每个索引都要写日志 + 可能页分裂
  · TRUNCATE 极便宜：只在 WAL 写一条「relfilenode 改为 N」元数据，磁盘上立即换文件
  · checkpoint 后第一次写每页都要 FPI（整页 8KB），所以单笔操作的 WAL 量被显著放大
  · wal_compression 会压缩 FPI，对 OLTP 写密集场景能省 30%~70% 的 WAL 体积
"""
        )


if __name__ == "__main__":
    main()