"""01_btree_vs_seqscan.py —— B-Tree 索引 vs 全表扫描

基于 init.sql 生成的 ch6_orders（100 万行）：
  1) 无索引查单用户订单（Seq Scan）
  2) 加索引后同一查询（Index Scan）
  3) 范围查询，对比 `(user_id)` 与 `(user_id, created_at)` 复合索引
"""
from __future__ import annotations

from _common import connect, explain_analyze, section, summarize_plan, time_query


def run(cur) -> None:
    # 清理可能存在的索引，保证重复运行
    cur.execute("DROP INDEX IF EXISTS idx_ch6_orders_user")
    cur.execute("DROP INDEX IF EXISTS idx_ch6_orders_user_time")

    USER_ID = 42
    SQL_POINT = "SELECT COUNT(*) FROM ch6_orders WHERE user_id = %s"
    SQL_RANGE = """
        SELECT id, amount, created_at FROM ch6_orders
        WHERE user_id = %s AND created_at >= NOW() - INTERVAL '90 days'
        ORDER BY created_at DESC LIMIT 20
    """

    section("1. 无索引 —— 百万行 Seq Scan")
    t0 = time_query(cur, SQL_POINT, (USER_ID,))
    p = summarize_plan(explain_analyze(cur, SQL_POINT, (USER_ID,)))
    print(f"  单点等值 COUNT  : best={t0:.2f}ms  node={p['node']}  "
          f"cost={p['total_cost']}  actual={p['actual_ms']}ms  "
          f"buffers hit={p['shared_hit']} read={p['shared_read']}")

    t1 = time_query(cur, SQL_RANGE, (USER_ID,))
    p = summarize_plan(explain_analyze(cur, SQL_RANGE, (USER_ID,)))
    print(f"  范围+排序  : best={t1:.2f}ms  node={p['node']}  "
          f"cost={p['total_cost']}  actual={p['actual_ms']}ms")

    section("2. 加 B-Tree 单列索引 (user_id)")
    cur.execute("CREATE INDEX idx_ch6_orders_user ON ch6_orders (user_id)")
    cur.execute("ANALYZE ch6_orders")

    t2 = time_query(cur, SQL_POINT, (USER_ID,))
    p = summarize_plan(explain_analyze(cur, SQL_POINT, (USER_ID,)))
    print(f"  单点等值 COUNT  : best={t2:.2f}ms  node={p['node']}  "
          f"cost={p['total_cost']}  actual={p['actual_ms']}ms  "
          f"buffers hit={p['shared_hit']} read={p['shared_read']}")
    print(f"  >>> 加速比: {(t0/t2):.1f}x")

    section("3. 加复合索引 (user_id, created_at DESC) —— 匹配范围+排序")
    cur.execute("DROP INDEX idx_ch6_orders_user")
    cur.execute("CREATE INDEX idx_ch6_orders_user_time "
                "ON ch6_orders (user_id, created_at DESC)")
    cur.execute("ANALYZE ch6_orders")

    t3 = time_query(cur, SQL_RANGE, (USER_ID,))
    p = summarize_plan(explain_analyze(cur, SQL_RANGE, (USER_ID,)))
    print(f"  范围+排序  : best={t3:.2f}ms  node={p['node']}  "
          f"cost={p['total_cost']}  actual={p['actual_ms']}ms")
    print(f"  >>> 对比无索引加速比: {(t1/t3):.1f}x")

    section("4. INCLUDE 实现 Index Only Scan")
    cur.execute("DROP INDEX idx_ch6_orders_user_time")
    cur.execute(
        "CREATE INDEX idx_ch6_orders_cover ON ch6_orders (user_id) "
        "INCLUDE (amount, status)"
    )
    cur.execute("ANALYZE ch6_orders")
    cur.execute("VACUUM ch6_orders")   # 让 VM 标记 all-visible，才能走 IOS

    SQL_COVER = (
        "SELECT amount, status FROM ch6_orders WHERE user_id = %s LIMIT 100"
    )
    p = summarize_plan(explain_analyze(cur, SQL_COVER, (USER_ID,)))
    print(f"  覆盖查询节点: {p['node']}  actual={p['actual_ms']}ms")
    print("  期望看到 'Index Only Scan'；若看到 'Index Scan'，说明 VM 尚未标记。")

    cur.execute("DROP INDEX idx_ch6_orders_cover")


def main() -> None:
    with connect() as conn:
        with conn.cursor() as cur:
            run(cur)
            conn.commit()


if __name__ == "__main__":
    main()