"""
02_hot_update.py —— HOT 更新 vs 非 HOT 更新对比

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

要点：
    ch8_hot_demo 表：
        - score 列没有索引   -> UPDATE score 可以 HOT
        - nickname 列有索引 -> UPDATE nickname 不能 HOT

    跑大量更新后对比：
        n_tup_upd / n_tup_hot_upd
        表大小 / 索引大小
"""
from __future__ import annotations

import psycopg

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


def reset_table() -> None:
    with psycopg.connect(DSN, autocommit=True) as conn:
        conn.execute("TRUNCATE ch8_hot_demo")
        conn.execute(
            "INSERT INTO ch8_hot_demo (id, nickname, score) "
            "SELECT g, 'user_' || g, g * 10 "
            "FROM generate_series(1, 1000) g"
        )
        conn.execute("VACUUM ch8_hot_demo")
        conn.execute("ANALYZE ch8_hot_demo")


def stats() -> tuple[int, int, str, str]:
    with psycopg.connect(DSN, autocommit=True) as conn:
        cur = conn.execute("""
            SELECT n_tup_upd, n_tup_hot_upd
            FROM pg_stat_user_tables
            WHERE relname='ch8_hot_demo'
        """)
        upd, hot = cur.fetchone() or (0, 0)
        cur = conn.execute(
            "SELECT pg_size_pretty(pg_relation_size('ch8_hot_demo')), "
            "pg_size_pretty(pg_relation_size('idx_ch8_hot_demo_nickname'))"
        )
        ts, idxs = cur.fetchone()
    return upd or 0, hot or 0, ts, idxs


def reset_stats() -> None:
    with psycopg.connect(DSN, autocommit=True) as conn:
        conn.execute("SELECT pg_stat_reset_single_table_counters("
                     "'ch8_hot_demo'::regclass)")


def run_updates(n: int, on_indexed_col: bool) -> None:
    sql = ("UPDATE ch8_hot_demo SET nickname = nickname || '+' "
           "WHERE id = %s") if on_indexed_col else (
          "UPDATE ch8_hot_demo SET score = score + 1 WHERE id = %s")
    with psycopg.connect(DSN, autocommit=True) as conn:
        for _ in range(n):
            for i in range(1, 11):
                conn.execute(sql, (i,))


def main() -> None:
    print("=== A. UPDATE 「无索引列 score」（应大量 HOT） ===")
    reset_table()
    reset_stats()
    run_updates(50, on_indexed_col=False)
    upd, hot, ts, idxs = stats()
    print(f"  n_tup_upd = {upd}, n_tup_hot_upd = {hot}, "
          f"HOT 比例 = {100*hot/max(1,upd):.1f}%")
    print(f"  表大小 = {ts}, 索引大小 = {idxs}")

    print("\n=== B. UPDATE 「有索引列 nickname」（应几乎无 HOT） ===")
    reset_table()
    reset_stats()
    run_updates(50, on_indexed_col=True)
    upd, hot, ts, idxs = stats()
    print(f"  n_tup_upd = {upd}, n_tup_hot_upd = {hot}, "
          f"HOT 比例 = {100*hot/max(1,upd):.1f}%")
    print(f"  表大小 = {ts}, 索引大小 = {idxs}")

    print("\n=== C. 在 B 之后再 VACUUM 看死元组与索引项 ===")
    with psycopg.connect(DSN, autocommit=True) as conn:
        conn.execute("VACUUM (VERBOSE) ch8_hot_demo")


if __name__ == "__main__":
    main()