"""
03_toast_demo.py
----------------
插入超大 text 列触发 TOAST：
  · 找到主表的 reltoastrelid（pg_toast.pg_toast_<oid>）
  · 比较主表 / TOAST 表的物理大小
  · 直接 SELECT pg_toast 表，看到一行被切成多个 chunk
  · 切换 STORAGE 策略（EXTENDED → EXTERNAL → MAIN → PLAIN），
    重新插入后看大小如何变化

要求：表 ch9_big_doc 已存在（init.sql 已创建）。
"""

import psycopg

CONN_STR = "host=127.0.0.1 port=5432 dbname=learn_pg user=postgres"
PAYLOAD_SIZE = 1024 * 1024  # 1MB


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


def main() -> None:
    with psycopg.connect(CONN_STR, autocommit=True) as conn, conn.cursor() as cur:
        cur.execute("TRUNCATE ch9_big_doc;")

        # ---------- 1. 找到 TOAST 表 ----------
        banner("1. 主表与 TOAST 表的对应关系")
        cur.execute(
            """
            SELECT c.oid       AS heap_oid,
                   c.relname,
                   t.oid       AS toast_oid,
                   t.relname   AS toast_relname,
                   pg_relation_filepath(t.oid) AS toast_path
            FROM pg_class c
            JOIN pg_class t ON c.reltoastrelid = t.oid
            WHERE c.relname = 'ch9_big_doc';
            """
        )
        heap_oid, relname, toast_oid, toast_name, toast_path = cur.fetchone()
        print(f"主表  : {relname:<10} oid={heap_oid}")
        print(f"TOAST : pg_toast.{toast_name}  oid={toast_oid}")
        print(f"路径  : {toast_path}")

        # ---------- 2. 插入 1MB 数据，观察 TOAST 表膨胀 ----------
        banner(f"2. 插入 {PAYLOAD_SIZE/1024:.0f}KB 字符串后两表大小对比")
        # 用可压缩的内容（重复字符）→ pglz 能压得很厉害
        cur.execute(
            "INSERT INTO ch9_big_doc (id, body) VALUES (%s, %s);",
            (1, "A" * PAYLOAD_SIZE),
        )
        # 用不易压缩的随机字节 → 几乎全部进 toast chunk
        cur.execute(
            "INSERT INTO ch9_big_doc (id, body) VALUES (%s, encode(gen_random_bytes(%s), 'hex'));",
            (2, PAYLOAD_SIZE // 2),
        )

        cur.execute(
            f"""
            SELECT
              pg_size_pretty(pg_relation_size('ch9_big_doc'))                AS heap_size,
              pg_size_pretty(pg_relation_size('pg_toast.{toast_name}'))  AS toast_size,
              pg_size_pretty(pg_total_relation_size('ch9_big_doc'))          AS total_size;
            """
        )
        h, t, total = cur.fetchone()
        print(f"  heap (ch9_big_doc)        = {h}")
        print(f"  toast ({toast_name}) = {t}")
        print(f"  total (heap+toast+ix) = {total}")

        # ---------- 3. 看 TOAST chunks ----------
        banner("3. TOAST 表中的 chunk 分布")
        cur.execute(
            f"""
            SELECT chunk_id,
                   COUNT(*)              AS chunk_count,
                   pg_size_pretty(SUM(length(chunk_data))) AS total_bytes
            FROM pg_toast.{toast_name}
            GROUP BY chunk_id
            ORDER BY chunk_id;
            """
        )
        print(f"  {'chunk_id':>10} {'#chunks':>10}  {'sum_bytes':>12}")
        for r in cur.fetchall():
            print(f"  {r[0]:>10} {r[1]:>10}  {r[2]:>12}")

        cur.execute(
            f"""
            SELECT chunk_id, chunk_seq, length(chunk_data) AS chunk_len
            FROM pg_toast.{toast_name}
            ORDER BY chunk_id, chunk_seq
            LIMIT 5;
            """
        )
        print("\n  前 5 片明细：")
        print(f"  {'chunk_id':>10} {'chunk_seq':>10}  chunk_len")
        for r in cur.fetchall():
            print(f"  {r[0]:>10} {r[1]:>10}  {r[2]}")
        print("\n  → 每片约 1996 字节（接近 BLCKSZ/4）")

        # ---------- 4. 切换 STORAGE 策略对比 ----------
        banner("4. STORAGE 策略影响（重写后看大小变化）")

        def measure(label: str) -> None:
            cur.execute("VACUUM FULL ch9_big_doc;")  # 重写以应用新策略
            cur.execute(
                f"""
                SELECT pg_size_pretty(pg_relation_size('ch9_big_doc')),
                       pg_size_pretty(pg_relation_size('pg_toast.{toast_name}'));
                """
            )
            h2, t2 = cur.fetchone()
            print(f"  {label:<30} heap={h2:<10}  toast={t2}")

        for storage in ("EXTENDED", "EXTERNAL", "MAIN", "PLAIN"):
            try:
                cur.execute(
                    f"ALTER TABLE ch9_big_doc ALTER COLUMN body SET STORAGE {storage};"
                )
                measure(f"STORAGE = {storage}")
            except psycopg.errors.ProgramLimitExceeded as e:
                # PLAIN/MAIN 可能因为单行 > 页大小而失败
                print(f"  STORAGE = {storage:<10} → 失败：{e}")
                # 回滚一下
                cur.execute("DELETE FROM ch9_big_doc;")
                cur.execute(
                    "INSERT INTO ch9_big_doc VALUES (1, repeat('A', 100));"
                )

        # ---------- 5. 关闭压缩看差异 ----------
        banner("5. 压缩算法（仅 PG14+）")
        try:
            cur.execute("SHOW default_toast_compression;")
            algo = cur.fetchone()[0]
            print(f"  当前默认 TOAST 压缩算法 = {algo}")
            print("  可选值：pglz（默认）/ lz4（速度快）")
            print("  改单列：ALTER TABLE t ALTER COLUMN c SET COMPRESSION lz4;")
        except psycopg.Error:
            print("  PG13 及以下：只有 pglz 一种压缩算法")


if __name__ == "__main__":
    main()