"""05_explain_diff.py —— 第 17 章配套代码 #5

用途
    自动比较「加索引前 vs 加索引后」同一条 SQL 的 EXPLAIN (ANALYZE, BUFFERS) 结果，
    凸显「Seq Scan → Index Scan」的飞跃。

测试场景
    SQL: SELECT * FROM ch17_perf_orders WHERE user_id = ? AND status = 1 ORDER BY created_at DESC LIMIT 20
    1) 不带索引：必然 Seq Scan + 在 WHERE 处过滤
    2) 加复合索引 (user_id, status, created_at DESC) 后：Index Scan + Limit 提前停

预期输出
    --- BEFORE ---
    Limit  ...
      ->  Sort  ...
            ->  Seq Scan on ch17_perf_orders  (rows=...) (cost=...)
                  Filter: ((user_id = N) AND (status = 1))
                  Rows Removed by Filter: 999800
                  Buffers: shared hit=12345 read=6789
    Planning Time: ...
    Execution Time: 380.2 ms

    --- AFTER  ---
    Limit  (...)
      ->  Index Scan using idx_perf_orders_us_st_ct on ch17_perf_orders
            Index Cond: ((user_id = N) AND (status = 1))
            Buffers: shared hit=23
    Planning Time: ...
    Execution Time: 0.6 ms
"""

from __future__ import annotations

import random
import sys
import textwrap

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"
INDEX_NAME = "idx_perf_orders_us_st_ct"
TARGET_SQL = """
SELECT id, user_id, product_id, status, amount, created_at
FROM ch17_perf_orders
WHERE user_id = %s AND status = 1
ORDER BY created_at DESC
LIMIT 20
"""


def explain(cur, sql: str, params: tuple) -> str:
    cur.execute("EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) " + sql, params)
    return "\n".join(r[0] for r in cur.fetchall())


def line(title: str) -> str:
    return "-" * 30 + f"  {title}  " + "-" * 30


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:
        # 让结果稳定：选一个真的存在的 user_id
        cur.execute("SELECT user_id FROM ch17_perf_orders LIMIT 1")
        uid = cur.fetchone()[0]
        print(f"测试 SQL：{textwrap.dedent(TARGET_SQL).strip()}")
        print(f"参数 user_id = {uid}\n")

        # 0. 确保索引被删干净
        cur.execute(f"DROP INDEX IF EXISTS {INDEX_NAME}")
        cur.execute("ANALYZE ch17_perf_orders")

        # 1. BEFORE
        print(line("BEFORE: 没有索引"))
        plan_before = explain(cur, TARGET_SQL, (uid,))
        print(plan_before)

        # 2. CREATE INDEX
        print()
        print(line(f"CREATE INDEX {INDEX_NAME}"))
        cur.execute(
            f"CREATE INDEX {INDEX_NAME} "
            "ON ch17_perf_orders (user_id, status, created_at DESC)"
        )
        cur.execute("ANALYZE ch17_perf_orders")
        print(f"已创建复合索引 (user_id, status, created_at DESC)")

        # 3. AFTER
        print()
        print(line("AFTER: 有索引"))
        plan_after = explain(cur, TARGET_SQL, (uid,))
        print(plan_after)

        # 4. 提取 Execution Time 数值做对比
        def exec_time(plan: str) -> float | None:
            for ln in plan.splitlines():
                if "Execution Time" in ln:
                    try:
                        return float(ln.split(":")[-1].replace("ms", "").strip())
                    except ValueError:
                        return None
            return None

        before, after = exec_time(plan_before), exec_time(plan_after)
        if before and after:
            print()
            print(line("对比结论"))
            speedup = before / after
            print(f"  Execution Time:  {before:.2f} ms  →  {after:.2f} ms")
            print(f"  提速：约 {speedup:,.0f} 倍")

        # 5. 清理（可选，注释掉则保留索引方便后续 demo）
        # cur.execute(f"DROP INDEX IF EXISTS {INDEX_NAME}")


if __name__ == "__main__":
    main()