"""03_keyset_pagination.py —— 第 17 章配套代码 #3

用途
    对比「OFFSET 巨大值的传统分页」与「键集分页 keyset pagination」的真实性能。
    OFFSET 越大越慢，因为 PG 必须先扫过前 OFFSET 行再丢弃。
    keyset 永远只扫一个索引区间，时间复杂度与 page 大小成正比，与「翻到第几页」无关。

依赖
    需要先运行 init.sql 准备 ch17_sorted_events 表（约 50w 行 + 复合索引）。

预期输出（示意，机器不同数值不同）
    OFFSET 0      LIMIT 50  ->   1.2 ms
    OFFSET 100000 LIMIT 50  ->  98.3 ms
    OFFSET 400000 LIMIT 50  -> 392.7 ms
    keyset 第 1 页           ->   1.0 ms
    keyset 第 2000 页（条件下沉）->   1.1 ms
"""

from __future__ import annotations

import sys
import time

try:
    import psycopg
except ImportError:
    sys.exit('请先安装 psycopg v3：pip install "psycopg[binary]>=3.1"')


CONN_INFO = "host=127.0.0.1 port=5432 dbname=learn_pg user=postgres"
PAGE_SIZE = 50


def timeit(cur, sql: str, params: tuple = ()) -> tuple[float, list]:
    """同一条 SQL 跑 3 次取最小值，去掉冷启动影响。"""
    best = float("inf")
    rows: list = []
    for _ in range(3):
        t0 = time.perf_counter()
        cur.execute(sql, params)
        rows = cur.fetchall()
        cost = (time.perf_counter() - t0) * 1000.0
        best = min(best, cost)
    return best, rows


def main() -> None:
    try:
        conn = psycopg.connect(CONN_INFO, autocommit=True)
    except psycopg.OperationalError as exc:
        sys.exit(f"连接失败：{exc}")

    with conn, conn.cursor() as cur:
        cur.execute("SELECT count(*) FROM ch17_sorted_events")
        total = cur.fetchone()[0]
        print(f"ch17_sorted_events 总行数：{total:,}")
        print(f"分页大小：{PAGE_SIZE} 行")

        # ---------- 一、OFFSET 分页 ----------
        print("\n[A] 传统 OFFSET 分页")
        print(f"{'OFFSET':>10} {'rows':>6}   {'耗时(ms)':>10}")
        for off in (0, 1_000, 10_000, 100_000, 400_000):
            cost, rows = timeit(
                cur,
                """
                SELECT id, happened_at, payload
                FROM ch17_sorted_events
                ORDER BY happened_at DESC, id DESC
                OFFSET %s LIMIT %s
                """,
                (off, PAGE_SIZE),
            )
            print(f"{off:>10,} {len(rows):>6}   {cost:>10.2f}")

        # ---------- 二、键集分页 ----------
        print("\n[B] 键集分页（keyset pagination）")
        print("    每次记住「上一页最后一行的 (happened_at, id)」")
        print("    新一页：WHERE (happened_at, id) < (?, ?) ORDER BY ... LIMIT N")
        print(f"{'第几页':>6} {'rows':>6}   {'耗时(ms)':>10}")

        # 模拟翻 5 页（每页都很快），再「跳到第 8000 页」也依然快
        last_ts, last_id = None, None
        for page in range(1, 6):
            if last_ts is None:
                cost, rows = timeit(
                    cur,
                    """
                    SELECT id, happened_at, payload
                    FROM ch17_sorted_events
                    ORDER BY happened_at DESC, id DESC
                    LIMIT %s
                    """,
                    (PAGE_SIZE,),
                )
            else:
                cost, rows = timeit(
                    cur,
                    """
                    SELECT id, happened_at, payload
                    FROM ch17_sorted_events
                    WHERE (happened_at, id) < (%s, %s)
                    ORDER BY happened_at DESC, id DESC
                    LIMIT %s
                    """,
                    (last_ts, last_id, PAGE_SIZE),
                )
            print(f"{page:>6} {len(rows):>6}   {cost:>10.2f}")
            if rows:
                last_id, last_ts = rows[-1][0], rows[-1][1]

        # 跳到「第 8000 页」（下沉条件 = 第 8000 页第 1 行的边界）
        # 我们用 OFFSET 取出第 8000*PAGE_SIZE 行的 (happened_at, id) 作为边界，
        # 然后用 keyset 取下一页 → 这一步即使在 50w 行规模下也不到几 ms。
        boundary_off = 8000 * PAGE_SIZE
        if boundary_off < total:
            cur.execute(
                """
                SELECT happened_at, id FROM ch17_sorted_events
                ORDER BY happened_at DESC, id DESC
                OFFSET %s LIMIT 1
                """,
                (boundary_off - 1,),
            )
            row = cur.fetchone()
            if row:
                bts, bid = row
                cost, rows = timeit(
                    cur,
                    """
                    SELECT id, happened_at, payload
                    FROM ch17_sorted_events
                    WHERE (happened_at, id) < (%s, %s)
                    ORDER BY happened_at DESC, id DESC
                    LIMIT %s
                    """,
                    (bts, bid, PAGE_SIZE),
                )
                print(f"{8001:>6} {len(rows):>6}   {cost:>10.2f}  (深翻一样快！)")

        # ---------- 三、EXPLAIN 直观对比 ----------
        print("\n[C] EXPLAIN 对比（看『扫了多少行』）")
        for label, sql, params in (
            ("OFFSET 400000",
             "EXPLAIN ANALYZE SELECT * FROM ch17_sorted_events "
             "ORDER BY happened_at DESC, id DESC OFFSET 400000 LIMIT 50",
             ()),
            ("keyset",
             "EXPLAIN ANALYZE SELECT * FROM ch17_sorted_events "
             "WHERE (happened_at, id) < (%s,%s) "
             "ORDER BY happened_at DESC, id DESC LIMIT 50",
             (last_ts, last_id)),
        ):
            print(f"\n--- {label} ---")
            cur.execute(sql, params)
            for r in cur.fetchall():
                print(" ", r[0])

        print("\n💡 总结：")
        print("   · OFFSET N LIMIT K → PG 必须扫过前 N 行才丢，深翻越来越慢")
        print("   · keyset 用复合索引精确定位边界，O(log N + K)，与页码无关")
        print("   · 注意排序键必须是「严格全序」，否则要带主键当 tie-breaker")


if __name__ == "__main__":
    main()