"""02_index_audit.py —— 第 17 章配套代码 #2

用途
    自动审计当前数据库的索引健康度，重点输出：
        ① 从未被使用过的索引（idx_scan = 0），是「光占空间还拖慢写入」的浪费
        ② 重复 / 冗余索引：索引 A 是索引 B 的前缀
        ③ 大表上没有索引的列：依据 pg_stats.n_distinct 推断「高基数列」候选
        ④ 索引膨胀粗略估算（精确膨胀需要 pgstattuple 扩展）

预期输出
    一张表格：
        schema | table         | index            | size   | scans | 建议
        public | ch17_perf_idx_demo | ch17_idx_demo_user    | 22 MB  |     0 | 未使用，建议 DROP
        public | ch17_perf_idx_demo | idx_demo_user_st | 31 MB  |     0 | 与 ch17_idx_demo_user_status_amt 前缀重复，建议合并
"""

from __future__ import annotations

import sys

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"


SQL_UNUSED = """
SELECT
    s.schemaname,
    s.relname  AS table_name,
    s.indexrelname AS index_name,
    pg_size_pretty(pg_relation_size(s.indexrelid)) AS size,
    s.idx_scan
FROM pg_stat_user_indexes s
JOIN pg_index i ON i.indexrelid = s.indexrelid
WHERE s.idx_scan = 0
  AND NOT i.indisunique         -- 唯一索引一般不能删（保业务约束）
  AND NOT i.indisprimary        -- 主键不能删
ORDER BY pg_relation_size(s.indexrelid) DESC
"""


SQL_DUPLICATE_PREFIX = """
WITH idx AS (
    SELECT
        n.nspname  AS schemaname,
        c.relname  AS table_name,
        ic.relname AS index_name,
        i.indexrelid,
        i.indrelid,
        ARRAY(
            SELECT pg_get_indexdef(i.indexrelid, k + 1, true)
            FROM generate_subscripts(i.indkey, 1) k
            ORDER BY k
        ) AS cols,
        pg_relation_size(i.indexrelid) AS bytes
    FROM pg_index i
    JOIN pg_class c   ON c.oid = i.indrelid
    JOIN pg_class ic  ON ic.oid = i.indexrelid
    JOIN pg_namespace n ON n.oid = c.relnamespace
    WHERE n.nspname NOT IN ('pg_catalog','information_schema')
      AND NOT i.indisprimary
      AND NOT i.indisunique
)
SELECT
    a.schemaname, a.table_name,
    a.index_name AS shorter_idx,
    b.index_name AS longer_idx,
    pg_size_pretty(a.bytes) AS shorter_size,
    pg_size_pretty(b.bytes) AS longer_size
FROM idx a
JOIN idx b
  ON a.indrelid = b.indrelid
 AND a.indexrelid <> b.indexrelid
 AND array_length(a.cols, 1) < array_length(b.cols, 1)
 AND a.cols = b.cols[1:array_length(a.cols,1)]
ORDER BY a.bytes DESC
"""


SQL_TABLE_NO_INDEX_CANDIDATE = """
SELECT
    s.schemaname,
    s.tablename,
    s.attname AS column_name,
    s.n_distinct,
    pg_size_pretty(pg_relation_size((s.schemaname||'.'||s.tablename)::regclass)) AS table_size
FROM pg_stats s
JOIN pg_class c ON c.relname = s.tablename
WHERE s.schemaname NOT IN ('pg_catalog','information_schema')
  AND pg_relation_size((s.schemaname||'.'||s.tablename)::regclass) > 50 * 1024 * 1024  -- > 50MB 才看
  AND s.n_distinct > 1000                                                             -- 高基数列
  AND NOT EXISTS (
      SELECT 1
      FROM pg_index i
      JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
      WHERE i.indrelid = (s.schemaname||'.'||s.tablename)::regclass
        AND a.attname  = s.attname
  )
ORDER BY pg_relation_size((s.schemaname||'.'||s.tablename)::regclass) DESC
"""


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


def print_table(headers: list[str], rows: list[tuple]) -> None:
    if not rows:
        print("  (无)")
        return
    widths = [max(len(str(h)), max((len(str(r[i])) for r in rows), default=0))
              for i, h in enumerate(headers)]
    print("  " + "  ".join(f"{h:<{widths[i]}}" for i, h in enumerate(headers)))
    print("  " + "  ".join("-" * w for w in widths))
    for r in rows:
        print("  " + "  ".join(f"{str(r[i]):<{widths[i]}}" for i in range(len(headers))))


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:
        banner("① 从未被使用过的索引（idx_scan = 0）")
        cur.execute(SQL_UNUSED)
        print_table(["schema", "table", "index", "size", "scans"], cur.fetchall())

        banner("② 前缀重复 / 冗余索引")
        cur.execute(SQL_DUPLICATE_PREFIX)
        print_table(
            ["schema", "table", "shorter_idx", "longer_idx", "shorter_size", "longer_size"],
            cur.fetchall(),
        )

        banner("③ 大表上「高基数 + 无索引」的列（建索引候选）")
        cur.execute(SQL_TABLE_NO_INDEX_CANDIDATE)
        print_table(
            ["schema", "table", "column", "n_distinct", "table_size"],
            cur.fetchall(),
        )

        banner("④ 各表索引/数据 体积比（看膨胀）")
        cur.execute("""
            SELECT
                schemaname,
                relname,
                pg_size_pretty(pg_relation_size(relid))                              AS table_size,
                pg_size_pretty(pg_indexes_size(relid))                               AS idx_size,
                ROUND(pg_indexes_size(relid)::NUMERIC
                      / NULLIF(pg_relation_size(relid),0), 2)                       AS idx_to_table_ratio
            FROM pg_stat_user_tables
            WHERE pg_relation_size(relid) > 10*1024*1024
            ORDER BY pg_indexes_size(relid) DESC
            LIMIT 20
        """)
        print_table(["schema", "table", "table_size", "idx_size", "ratio"], cur.fetchall())

        print("\n💡 经验值：")
        print("   · idx_to_table_ratio > 1   → 索引比表还大，多半有冗余索引")
        print("   · idx_scan = 0 持续 N 周  → 该索引在生产里真的没人用，可以 DROP")
        print("   · 同 indrelid 下，A 的列是 B 的前缀 → A 几乎可以被 B 替代")


if __name__ == "__main__":
    main()