"""公共工具：psycopg v3 连接 + 计时 + 美化 EXPLAIN。

依赖：pip install "psycopg[binary]"
"""
from __future__ import annotations

import json
import os
import time
from contextlib import contextmanager
from typing import Iterable, Sequence

import psycopg


def conninfo() -> str:
    host = os.getenv("PGHOST", "127.0.0.1")
    port = os.getenv("PGPORT", "5432")
    db   = os.getenv("PGDATABASE", "learn_pg")
    user = os.getenv("PGUSER", "postgres")
    pwd  = os.getenv("PGPASSWORD", "")
    parts = [f"host={host}", f"port={port}", f"dbname={db}", f"user={user}"]
    if pwd:
        parts.append(f"password={pwd}")
    return " ".join(parts)


@contextmanager
def connect():
    with psycopg.connect(conninfo(), autocommit=False) as conn:
        yield conn


def section(title: str) -> None:
    bar = "=" * 76
    print("\n" + bar)
    print(f"  {title}")
    print(bar)


def print_table(rows: Sequence[Sequence], headers: Iterable[str]) -> None:
    headers = list(headers)
    str_rows = [[("" if v is None else str(v)) for v in r] for r in rows]
    widths = [len(h) for h in headers]
    for r in str_rows:
        for i, v in enumerate(r):
            if i < len(widths):
                widths[i] = max(widths[i], len(v))
    line = "+" + "+".join("-" * (w + 2) for w in widths) + "+"
    print(line)
    print("| " + " | ".join(h.ljust(widths[i]) for i, h in enumerate(headers)) + " |")
    print(line)
    for r in str_rows:
        print("| " + " | ".join(r[i].ljust(widths[i]) for i in range(len(widths))) + " |")
    print(line)


def time_query(cur, sql: str, params: tuple | None = None, runs: int = 3) -> float:
    """取 runs 次的最小耗时（毫秒），消除抖动。"""
    best = float("inf")
    for _ in range(runs):
        t0 = time.perf_counter()
        cur.execute(sql, params or ())
        cur.fetchall()
        dt = (time.perf_counter() - t0) * 1000
        best = min(best, dt)
    return best


def explain_analyze(cur, sql: str, params: tuple | None = None) -> dict:
    """执行 EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) 返回 plan 字典。"""
    cur.execute("EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) " + sql, params or ())
    return cur.fetchone()[0][0]


def summarize_plan(plan_root: dict) -> dict:
    """从 EXPLAIN JSON 提炼关键指标。"""
    top = plan_root["Plan"]
    return {
        "node": top["Node Type"],
        "total_cost": round(top["Total Cost"], 2),
        "actual_ms": round(top["Actual Total Time"], 3),
        "rows": top["Actual Rows"],
        "shared_hit": top.get("Shared Hit Blocks", 0),
        "shared_read": top.get("Shared Read Blocks", 0),
        "planning_ms": round(plan_root.get("Planning Time", 0), 3),
        "exec_ms": round(plan_root.get("Execution Time", 0), 3),
    }