"""05_pg_trgm_search.py —— 第 18 章配套代码 #5

用途
    用 pg_trgm 三元组扩展实现「商品名称模糊搜索」：
        ① LIKE '%关键词%' 在大表上一般是全表扫；建 GIN trigram 索引后能秒级返回
        ② similarity(a, b) 返回 0~1 的相似度，可用于「拼写纠错」
        ③ % 操作符 = similarity 阈值过滤

依赖
    · pg_trgm 已启用（init.sql 里已 CREATE EXTENSION pg_trgm）
    · ch18_goods 表已建好 + GIN trigram 索引

预期
    [无索引等价查询] LIKE '%queryword%'  -> Seq Scan ~XX ms
    [GIN trgm 索引]  LIKE '%queryword%'  -> Bitmap Index Scan ~Y ms
    Top 5 相似词:
       1) 0.84  iPhone 15 Pro Max ...
       2) 0.71  ...
"""

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"


def banner(t: str) -> None:
    print()
    print("=" * 80)
    print(f" {t}")
    print("=" * 80)


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 1 FROM pg_extension WHERE extname = 'pg_trgm'")
        if cur.fetchone() is None:
            sys.exit("❌ pg_trgm 未启用：请先 CREATE EXTENSION pg_trgm（init.sql 已包含）")

        cur.execute("SELECT count(*) FROM ch18_goods")
        total = cur.fetchone()[0]
        print(f"ch18_goods 表行数：{total:,}")

        # 1. 直接 LIKE，看 EXPLAIN 是否走 trgm 索引
        kw = "iPhone"
        banner(f"① LIKE '%{kw}%' + GIN trigram 索引")
        sql = "SELECT id, name FROM ch18_goods WHERE name LIKE %s LIMIT 10"
        cur.execute("EXPLAIN ANALYZE " + sql, (f"%{kw}%",))
        for r in cur.fetchall():
            print("  ", r[0])
        t0 = time.perf_counter()
        cur.execute(sql, (f"%{kw}%",))
        rows = cur.fetchall()
        print(f"  → {(time.perf_counter()-t0)*1000:.2f} ms，前 10 条：")
        for i, (gid, name) in enumerate(rows, 1):
            print(f"     {i}) #{gid} {name}")

        # 2. 临时禁用索引看「裸跑」对比
        banner(f"② 对照：禁用 trgm 索引后跑同样 LIKE")
        cur.execute("SET LOCAL enable_seqscan = on")
        cur.execute("SET LOCAL enable_indexscan = off")
        cur.execute("SET LOCAL enable_bitmapscan = off")
        cur.execute("EXPLAIN ANALYZE " + sql, (f"%{kw}%",))
        for r in cur.fetchall():
            print("  ", r[0])
        cur.execute("RESET enable_seqscan")
        cur.execute("RESET enable_indexscan")
        cur.execute("RESET enable_bitmapscan")

        # 3. similarity 操作符：模糊匹配 + 拼写纠错
        misspell = "iPhne 15 Pr"
        banner(f"③ similarity 拼写纠错：用户输入「{misspell}」")
        cur.execute(
            """
            SELECT name, similarity(name, %s) AS sim
            FROM ch18_goods
            WHERE name %% %s            -- '%' 操作符等价于 similarity > pg_trgm.similarity_threshold
            ORDER BY sim DESC
            LIMIT 5
            """,
            (misspell, misspell),
        )
        for i, (name, sim) in enumerate(cur.fetchall(), 1):
            print(f"   {i}) sim={sim:.3f}  {name}")

        # 4. 当前 similarity_threshold
        cur.execute("SHOW pg_trgm.similarity_threshold")
        thresh = cur.fetchone()[0]
        print(f"\n  当前 pg_trgm.similarity_threshold = {thresh}")
        print("  调整：SET pg_trgm.similarity_threshold = 0.2;  // 默认 0.3")

        # 5. <-> 距离操作符 + ORDER BY KNN
        banner("④ <-> 距离操作符（最相似 Top N，无需 WHERE）")
        cur.execute(
            """
            SELECT name, name <-> %s AS distance
            FROM ch18_goods
            ORDER BY name <-> %s
            LIMIT 5
            """,
            (misspell, misspell),
        )
        for i, (name, dist) in enumerate(cur.fetchall(), 1):
            print(f"   {i}) dist={dist:.3f}  {name}")

        print("\n💡 pg_trgm 速查：")
        print("   · 字符串切成连续 3 字符片段（trigram）做集合相似度")
        print("   · LIKE '%xxx%' / ILIKE → GIN trigram 索引")
        print("   · % 操作符 → similarity 阈值过滤")
        print("   · <-> 操作符 → 距离 = 1 - similarity，ORDER BY 拿 Top N")
        print("   · 中文也能用，但建议长度 ≥ 3 个字符")


if __name__ == "__main__":
    main()