"""04_explain_analyze.py —— 用 psycopg 抓 EXPLAIN JSON 并美化输出

功能：
  1) 递归遍历 EXPLAIN JSON 树
  2) 用 ANSI 颜色标注昂贵节点、估算偏差大的节点
  3) 高亮 Buffers 命中率低的节点
  4) 最终给出 3 条调优建议
"""
from __future__ import annotations

import json
import sys

from _common import connect, section


def color(s: str, c: str) -> str:
    codes = {
        "red": "31", "yellow": "33", "green": "32", "cyan": "36",
        "gray": "90", "bold": "1",
    }
    return f"\x1b[{codes[c]}m{s}\x1b[0m"


def render(node: dict, depth: int = 0, suggestions: list | None = None) -> None:
    """递归打印单个 Plan 节点。"""
    if suggestions is None:
        suggestions = []
    prefix = "  " * depth + ("└─ " if depth else "")

    nt = node["Node Type"]
    cost = node.get("Total Cost", 0)
    rows_est = node.get("Plan Rows", 0)
    rows_act = node.get("Actual Rows", 0)
    loops = node.get("Actual Loops", 1)
    ms = node.get("Actual Total Time", 0)
    hit = node.get("Shared Hit Blocks", 0)
    read = node.get("Shared Read Blocks", 0)

    # 估算偏差比（避免除 0）
    bias = None
    if rows_est > 0 and rows_act > 0:
        bias = rows_act / rows_est if rows_act >= rows_est else rows_est / rows_act

    head = f"{prefix}{color(nt, 'bold')}  cost={cost:.1f}"
    head += f"  est_rows={rows_est}  actual_rows={rows_act}"
    if loops > 1:
        head += f"  loops={loops}"
    head += f"  actual_ms={ms:.3f}"
    if hit or read:
        head += f"  buffers(hit={hit}, read={read})"
    print(head)

    # 附加信息（Filter、Index Cond、Sort Key 等）
    for key in ("Index Cond", "Recheck Cond", "Filter", "Hash Cond",
                "Sort Key", "Group Key", "Join Filter", "Merge Cond"):
        if key in node:
            print("  " * (depth + 1) + color(f"{key}: ", "gray") + str(node[key]))

    # 规则告警
    if nt == "Seq Scan" and rows_act > 10000:
        msg = f"大表 Seq Scan 返回 {rows_act} 行 → 考虑加索引"
        print("  " * (depth + 1) + color(f"⚠ {msg}", "red"))
        suggestions.append(msg)
    if bias and bias > 10:
        msg = (f"{nt} 估算 {rows_est} vs 实际 {rows_act}，"
               f"偏差 {bias:.1f}× → 跑 ANALYZE 或扩展统计")
        print("  " * (depth + 1) + color(f"⚠ {msg}", "yellow"))
        suggestions.append(msg)
    if read > 0 and hit + read > 100 and read > hit:
        msg = f"{nt} 缓冲读磁盘 {read} > hit {hit} → 缓存未命中，考虑 warm cache"
        print("  " * (depth + 1) + color(f"⚠ {msg}", "yellow"))
        suggestions.append(msg)

    for child in node.get("Plans", []) or []:
        render(child, depth + 1, suggestions)


def analyze(cur, sql: str, params: tuple | None = None) -> None:
    section(f"SQL: {sql.strip()[:80]}")
    cur.execute("EXPLAIN (ANALYZE, BUFFERS, VERBOSE, FORMAT JSON) " + sql, params or ())
    plan = cur.fetchone()[0][0]

    suggestions: list[str] = []
    render(plan["Plan"], 0, suggestions)

    print()
    print(f"Planning Time:  {plan.get('Planning Time', 0):.3f} ms")
    print(f"Execution Time: {plan.get('Execution Time', 0):.3f} ms")

    if suggestions:
        print("\n" + color("📌 调优建议", "cyan"))
        for i, s in enumerate(set(suggestions), 1):
            print(f"  {i}. {s}")
    else:
        print("\n" + color("✓ 未发现显著性能问题", "green"))


def main() -> None:
    demos = [
        # 1) 故意没建索引 → Seq Scan
        ("大表 Seq Scan 慢查询",
         "SELECT COUNT(*) FROM ch6_orders WHERE amount > 4800", None),

        # 2) 配合索引 → Index Scan
        ("复合条件：期望走 Index Scan",
         "SELECT id FROM ch6_orders WHERE user_id = %s AND created_at >= %s "
         "ORDER BY created_at DESC LIMIT 20",
         (42, "2024-01-01")),

        # 3) JOIN 聚合
        ("JOIN + GROUP BY 复杂计划",
         """
         SELECT u.city, COUNT(*) AS cnt
         FROM ch6_users u JOIN ch6_orders o ON o.user_id = u.id
         WHERE o.created_at >= NOW() - INTERVAL '30 days'
         GROUP BY u.city ORDER BY cnt DESC
         """, None),
    ]

    with connect() as conn:
        with conn.cursor() as cur:
            for title, sql, params in demos:
                try:
                    print("\n" + color(">>> " + title, "cyan"))
                    analyze(cur, sql, params)
                except Exception as e:
                    print(color(f"跳过（{type(e).__name__}: {e}）", "gray"))


if __name__ == "__main__":
    main()