"""02_postgis_nearby.py —— 第 18 章配套代码 #2

用途
    用 PostGIS 实现「附近 N 公里的商家」查询，并对比有/无 GiST 空间索引的耗时。
    场景：用户站在天安门 (116.397428, 39.90923)，找 3 公里以内的 POI。

依赖
    1. PostGIS：apt install postgresql-XX-postgis-3
    2. 已运行 init.sql 准备 ch18_poi 表（5000 行 POI）

预期输出
    [无索引] 全表扫            : XX ms, 找到 N 个
    [GiST 索引] ST_DWithin     : YY ms, 找到 N 个
    Top 5 最近 POI:
      1) POI #321  cafe        125.3 m
      ...
"""

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"
USER_LON, USER_LAT = 116.397428, 39.90923
RADIUS_M = 3000


def has_postgis(cur) -> bool:
    cur.execute("SELECT 1 FROM pg_extension WHERE extname = 'postgis'")
    return cur.fetchone() is not None


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:
        if not has_postgis(cur):
            sys.exit(
                "❌ PostGIS 未启用。请先：\n"
                "   apt install postgresql-XX-postgis-3 (XX = PG 主版本号)\n"
                "   psql -U postgres -d learn_pg -c 'CREATE EXTENSION postgis;'\n"
                "   然后重跑 init.sql 生成 ch18_poi 表。"
            )

        cur.execute("SELECT count(*) FROM ch18_poi")
        total = cur.fetchone()[0]
        print(f"POI 总数：{total}，用户位置 ({USER_LON}, {USER_LAT})，搜索半径 {RADIUS_M} 米\n")

        # 1. 不用空间索引：用 ST_Distance > radius 强制不走索引
        print("[1] 暴力距离过滤（不会走索引，因为函数表达式两边都是变量）")
        cur.execute("DROP INDEX IF EXISTS idx_poi_geom_tmp")  # 确保不影响
        sql_no_idx = """
            SELECT id, name, kind,
                   ST_Distance(geom, ST_MakePoint(%s, %s)::geography) AS dist_m
            FROM ch18_poi
            WHERE ST_Distance(geom, ST_MakePoint(%s, %s)::geography) <= %s
            ORDER BY dist_m
        """
        t0 = time.perf_counter()
        cur.execute(sql_no_idx, (USER_LON, USER_LAT, USER_LON, USER_LAT, RADIUS_M))
        no_idx_rows = cur.fetchall()
        cost_no_idx = (time.perf_counter() - t0) * 1000
        print(f"  → {cost_no_idx:.1f} ms，找到 {len(no_idx_rows)} 个 POI")

        # 2. 用 GiST 索引 + ST_DWithin（PostGIS 推荐姿势）
        print("\n[2] ST_DWithin + GiST 索引（推荐姿势）")
        sql_idx = """
            SELECT id, name, kind,
                   ST_Distance(geom, ST_MakePoint(%s, %s)::geography) AS dist_m
            FROM ch18_poi
            WHERE ST_DWithin(geom, ST_MakePoint(%s, %s)::geography, %s)
            ORDER BY geom <-> ST_MakePoint(%s, %s)::geography
            LIMIT 50
        """
        t0 = time.perf_counter()
        cur.execute(sql_idx, (USER_LON, USER_LAT, USER_LON, USER_LAT, RADIUS_M, USER_LON, USER_LAT))
        idx_rows = cur.fetchall()
        cost_idx = (time.perf_counter() - t0) * 1000
        print(f"  → {cost_idx:.1f} ms，返回前 50 个最近 POI")

        # 3. 看 EXPLAIN
        print("\n[3] EXPLAIN ANALYZE（验证索引是否生效）")
        cur.execute(
            "EXPLAIN (ANALYZE, BUFFERS) " + sql_idx,
            (USER_LON, USER_LAT, USER_LON, USER_LAT, RADIUS_M, USER_LON, USER_LAT),
        )
        for row in cur.fetchall():
            print("  ", row[0])

        # 4. Top 5
        print("\n[4] Top 5 最近 POI")
        for i, (pid, name, kind, dist) in enumerate(idx_rows[:5], 1):
            print(f"   {i}) #{pid:<6} {kind:<10} {name:<30} {dist:>7.1f} m")

        # 5. 关键 PostGIS 操作符 / 函数小抄
        print("\n💡 PostGIS 关键 API 速查：")
        print("   · ST_MakePoint(lon, lat)              → POINT 几何")
        print("   · ::geography                          → 转地理类型，距离单位为米")
        print("   · ST_DWithin(g1, g2, dist)             → 「g1 到 g2 距离 ≤ dist」（走 GiST）")
        print("   · ST_Distance(g1, g2)                  → 精确距离")
        print("   · g1 <-> g2                            → 「KNN 距离操作符」，配合 ORDER BY 用")
        print("   · ST_AsGeoJSON(g)                      → 输出 GeoJSON 给前端")
        print("   · CREATE INDEX ... USING gist (geom)  → 必建空间索引")


if __name__ == "__main__":
    main()