"""04_copy_vs_insert.py —— 第 17 章配套代码 #4

用途
    对比三种写入方式的吞吐：
        ① 单行 INSERT（最慢，每行一次网络往返）
        ② 多行 INSERT（一次发 1000 行，省网络往返）
        ③ COPY  FROM STDIN（PG 的「批量装载协议」，写入速度最快）
    输入数据：模拟 50000 行订单。

预期结果（量级，机器越好差距越大）
    单行 INSERT       :   50000 rows  /  6.5 s  →   7,700 rows/s
    多行 INSERT(1000) :   50000 rows  /  0.55 s →  91,000 rows/s
    COPY  FROM STDIN  :   50000 rows  /  0.12 s → 416,000 rows/s

结论
    生产里凡是「批量导入 / 数据迁移 / ETL」都应该用 COPY；
    REST API 里收到的数组写入也尽量攒批用多行 INSERT。
"""

from __future__ import annotations

import io
import random
import sys
import time
from datetime import datetime, timedelta, timezone

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"
ROWS = 50_000
BATCH = 1000


def gen_rows(n: int):
    base = datetime.now(tz=timezone.utc)
    for i in range(n):
        yield (
            random.randint(1, 50_000),
            random.randint(1, 1000),
            random.randint(0, 4),
            round(random.random() * 9999 + 1, 2),
            base - timedelta(seconds=i),
            f"perf_note_{i}",
        )


def setup(cur) -> None:
    cur.execute("DROP TABLE IF EXISTS ch17_perf_write_demo")
    cur.execute(
        """
        CREATE UNLOGGED TABLE ch17_perf_write_demo (
            id           BIGSERIAL PRIMARY KEY,
            user_id      BIGINT,
            product_id   BIGINT,
            status       SMALLINT,
            amount       NUMERIC(12,2),
            created_at   TIMESTAMPTZ,
            note         TEXT
        )
        """
    )


def truncate(cur) -> None:
    cur.execute("TRUNCATE ch17_perf_write_demo RESTART IDENTITY")


def bench_single_insert(cur, rows) -> float:
    truncate(cur)
    t0 = time.perf_counter()
    sql = ("INSERT INTO ch17_perf_write_demo "
           "(user_id,product_id,status,amount,created_at,note) "
           "VALUES (%s,%s,%s,%s,%s,%s)")
    for r in rows:
        cur.execute(sql, r)
    return time.perf_counter() - t0


def bench_multi_insert(cur, rows, batch: int = BATCH) -> float:
    truncate(cur)
    t0 = time.perf_counter()
    buf: list = []
    for r in rows:
        buf.append(r)
        if len(buf) >= batch:
            cur.executemany(
                "INSERT INTO ch17_perf_write_demo "
                "(user_id,product_id,status,amount,created_at,note) "
                "VALUES (%s,%s,%s,%s,%s,%s)",
                buf,
            )
            buf.clear()
    if buf:
        cur.executemany(
            "INSERT INTO ch17_perf_write_demo "
            "(user_id,product_id,status,amount,created_at,note) "
            "VALUES (%s,%s,%s,%s,%s,%s)",
            buf,
        )
    return time.perf_counter() - t0


def bench_copy(cur, rows) -> float:
    truncate(cur)
    t0 = time.perf_counter()
    with cur.copy(
        "COPY ch17_perf_write_demo "
        "(user_id,product_id,status,amount,created_at,note) FROM STDIN"
    ) as cp:
        for r in rows:
            cp.write_row(r)
    return time.perf_counter() - t0


def fmt(n_rows: int, sec: float) -> str:
    rps = n_rows / sec if sec > 0 else float("inf")
    return f"{n_rows:>7,} rows  /  {sec:6.2f} s  →  {rps:>10,.0f} rows/s"


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:
        setup(cur)

        print(f"基准：{ROWS:,} 行写入对比（表为 UNLOGGED 以排除 WAL 噪声）\n")

        rows1 = list(gen_rows(ROWS))
        cost = bench_single_insert(cur, rows1)
        print("① 单行 INSERT          :", fmt(ROWS, cost))

        rows2 = list(gen_rows(ROWS))
        cost = bench_multi_insert(cur, rows2, BATCH)
        print(f"② 批量 INSERT batch={BATCH} :", fmt(ROWS, cost))

        rows3 = list(gen_rows(ROWS))
        cost = bench_copy(cur, rows3)
        print("③ COPY FROM STDIN      :", fmt(ROWS, cost))

        print("\n💡 结论：")
        print("   · 单行 INSERT 慢 ≠ PG 慢，是「客户端往返 + 解析 + 计划」开销")
        print("   · 批量 INSERT 一次发 N 行，能把 RTT 摊薄成 1/N")
        print("   · COPY 走的是专门的二进制协议，跳过 SQL 解析阶段，最快")
        print("   · 如果要保 ACID 持久性，去掉 UNLOGGED 改成普通表，COPY 仍然显著快")


if __name__ == "__main__":
    main()