"""02_gin_jsonb.py —— GIN 索引在 JSONB 字段上的威力

基于 init.sql 的 ch6_docs（10 万行，每行是一个 JSONB）：
  1) 无索引查「有 vip=true 标签的文档」
  2) 加 GIN（jsonb_path_ops）索引后再查
  3) 对比 jsonb_ops vs jsonb_path_ops 的索引体积
  4) 演示 @> / ? / ?& 运算符的执行计划
"""
from __future__ import annotations

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


QUERIES = {
    "vip=true 筛选":        ("SELECT COUNT(*) FROM ch6_docs WHERE tags @> %s",
                             ('{"vip": true}',)),
    "country=JP 筛选":      ("SELECT COUNT(*) FROM ch6_docs WHERE tags @> %s",
                             ('{"country": "JP"}',)),
    "level=3 筛选":         ("SELECT COUNT(*) FROM ch6_docs WHERE tags @> %s",
                             ('{"level": 3}',)),
}


def run_batch(cur, title: str) -> None:
    section(title)
    for name, (sql, params) in QUERIES.items():
        t = time_query(cur, sql, params, runs=3)
        plan = summarize_plan(explain_analyze(cur, sql, params))
        print(f"  {name:22s}  node={plan['node']:20s}  "
              f"actual={plan['actual_ms']:.3f}ms  rows={plan['rows']}  "
              f"best={t:.2f}ms")


def index_size(cur, idx_name: str) -> str:
    cur.execute("SELECT pg_size_pretty(pg_relation_size(%s))", (idx_name,))
    (sz,) = cur.fetchone()
    return sz


def main() -> None:
    with connect() as conn:
        with conn.cursor() as cur:
            cur.execute("DROP INDEX IF EXISTS idx_ch6_docs_gin_ops")
            cur.execute("DROP INDEX IF EXISTS idx_ch6_docs_gin_path")

            run_batch(cur, "1. 无索引 —— 全表 Seq Scan")

            section("2. 建默认 GIN (jsonb_ops)")
            cur.execute("CREATE INDEX idx_ch6_docs_gin_ops ON ch6_docs USING GIN (tags)")
            cur.execute("ANALYZE ch6_docs")
            print(f"  索引体积: {index_size(cur, 'idx_ch6_docs_gin_ops')}")
            run_batch(cur, "   查询表现")

            cur.execute("DROP INDEX idx_ch6_docs_gin_ops")

            section("3. 建 GIN (jsonb_path_ops) —— 更小更快但只支持 @>")
            cur.execute(
                "CREATE INDEX idx_ch6_docs_gin_path ON ch6_docs "
                "USING GIN (tags jsonb_path_ops)"
            )
            cur.execute("ANALYZE ch6_docs")
            print(f"  索引体积: {index_size(cur, 'idx_ch6_docs_gin_path')}")
            run_batch(cur, "   查询表现")

            section("4. 顶层 key 存在查询 `tags ? 'vip'`（jsonb_path_ops 走不了）")
            sql = "SELECT COUNT(*) FROM ch6_docs WHERE tags ? %s"
            plan = summarize_plan(explain_analyze(cur, sql, ('vip',)))
            print(f"  节点: {plan['node']}  actual={plan['actual_ms']}ms")
            print("  说明：jsonb_path_ops 不支持 `?` 运算符，此时只能 Seq Scan。")
            print("  若业务需要 `?`，换回默认 jsonb_ops。")

            cur.execute("DROP INDEX idx_ch6_docs_gin_path")
            conn.commit()


if __name__ == "__main__":
    main()