"""basic_query.py —— 第 2 章配套代码 · 图书管理系统

用途
    用 psycopg v3 演示连接、查询、参数化查询、INSERT ... RETURNING、
    事务（BEGIN/COMMIT/ROLLBACK）、批量执行（executemany）、
    高效批量加载（COPY）等 PG 常用编程姿势。

前置
    1. 已经按第 2 章 init.sql 建好 ch2_books / ch2_authors / ch2_categories 三张表：
           psql -h 127.0.0.1 -U postgres -d learn_pg -f ../init.sql
    2. 安装 psycopg v3：
           pip install "psycopg[binary]>=3.1"

运行
    python basic_query.py
"""

from __future__ import annotations

import io
import sys
from decimal import Decimal

try:
    import psycopg
    from psycopg import sql
    from psycopg.rows import dict_row
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"


def section(title: str) -> None:
    line = "─" * 70
    print(f"\n{line}\n  {title}\n{line}")


# ------------------------------------------------------------------ #
#  Demo 1：基础查询 + 参数化（防 SQL 注入）
# ------------------------------------------------------------------ #
def demo_select_basic(conn: psycopg.Connection) -> None:
    section("Demo 1 · 基础查询 + 参数化查询")

    with conn.cursor(row_factory=dict_row) as cur:
        cur.execute("SELECT count(*) AS n FROM ch2_books")
        print("ch2_books 总条数 :", cur.fetchone()["n"])

        # 参数化查询：%s 占位符（不是 Python 的 % 格式化！）
        category = "科幻"
        min_stock = 1
        cur.execute(
            """
            SELECT b.id, b.title, a.name AS author, b.price, b.stock
            FROM ch2_books b
            JOIN ch2_authors    a ON a.id = b.author_id
            JOIN ch2_categories c ON c.id = b.category_id
            WHERE c.name = %s AND b.stock >= %s
            ORDER BY b.price DESC
            FETCH FIRST 5 ROWS ONLY
            """,
            (category, min_stock),
        )
        print(f"分类 = {category}, 库存 >= {min_stock} 的 Top5：")
        for row in cur.fetchall():
            print(
                f"  #{row['id']:<3}  {row['title']:<30}  "
                f"author={row['author']:<10}  ¥{row['price']}  stock={row['stock']}"
            )


# ------------------------------------------------------------------ #
#  Demo 2：INSERT ... RETURNING（PG 特色）
# ------------------------------------------------------------------ #
def demo_insert_returning(conn: psycopg.Connection) -> int:
    section("Demo 2 · INSERT ... RETURNING 一次拿到主键 + 时间戳")

    with conn.cursor() as cur:
        cur.execute(
            """
            INSERT INTO ch2_books (title, author_id, category_id, isbn, published_at, stock, price)
            VALUES (%s, %s, %s, %s, %s, %s, %s)
            RETURNING id, created_at
            """,
            (
                "psycopg 实战手册",
                5,                     # Knuth 占个位
                3,                     # 技术分类
                "9999999999999",
                "2026-04-01",
                100,
                Decimal("59.90"),
            ),
        )
        new_id, created_at = cur.fetchone()
        print(f"  新书已插入：id={new_id}, created_at={created_at}")

    conn.commit()
    return new_id


# ------------------------------------------------------------------ #
#  Demo 3：事务 —— BEGIN / COMMIT / ROLLBACK
# ------------------------------------------------------------------ #
def demo_transaction(conn: psycopg.Connection, book_id: int) -> None:
    section("Demo 3 · 事务的 ROLLBACK 保护数据安全")

    with conn.cursor() as cur:
        cur.execute("SELECT stock FROM ch2_books WHERE id = %s", (book_id,))
        before = cur.fetchone()[0]
        print(f"  事务前 stock = {before}")

    # psycopg v3 默认就是事务模式（autocommit=False），
    # with conn.transaction() 显式开启嵌套事务 / 保存点。
    try:
        with conn.transaction():
            with conn.cursor() as cur:
                cur.execute(
                    "UPDATE ch2_books SET stock = stock + 999 WHERE id = %s", (book_id,)
                )
            raise RuntimeError("人为抛出错误，触发回滚")
    except RuntimeError as exc:
        print(f"  捕获异常并回滚：{exc}")

    with conn.cursor() as cur:
        cur.execute("SELECT stock FROM ch2_books WHERE id = %s", (book_id,))
        after = cur.fetchone()[0]
        print(f"  事务后 stock = {after}（与事务前一致 → 回滚成功）")


# ------------------------------------------------------------------ #
#  Demo 4：executemany —— 批量插入
# ------------------------------------------------------------------ #
def demo_executemany(conn: psycopg.Connection) -> None:
    section("Demo 4 · executemany 批量插入 5 个新作者")

    rows = [
        ("Sample Author A", "测试国", 2000),
        ("Sample Author B", "测试国", 2001),
        ("Sample Author C", "测试国", 2002),
        ("Sample Author D", "测试国", 2003),
        ("Sample Author E", "测试国", 2004),
    ]

    with conn.cursor() as cur:
        cur.executemany(
            "INSERT INTO ch2_authors(name, country, born_year) VALUES (%s, %s, %s)",
            rows,
        )
        print(f"  已插入 {cur.rowcount} 行")

        cur.execute(
            "DELETE FROM ch2_authors WHERE name LIKE 'Sample Author%' RETURNING id"
        )
        print(f"  已清理 {cur.rowcount} 行（避免污染数据）")
    conn.commit()


# ------------------------------------------------------------------ #
#  Demo 5：COPY —— PG 最快的批量加载
# ------------------------------------------------------------------ #
def demo_copy(conn: psycopg.Connection) -> None:
    section("Demo 5 · COPY FROM STDIN 高速批量加载（10x ~ 100x 快于 INSERT）")

    csv_data = io.StringIO()
    for i in range(20):
        csv_data.write(f"copy_demo_author_{i}\t测试国\t{1900 + i}\n")
    csv_data.seek(0)

    with conn.cursor() as cur:
        with cur.copy(
            "COPY ch2_authors(name, country, born_year) FROM STDIN WITH (FORMAT text)"
        ) as cp:
            cp.write(csv_data.read())

        cur.execute(
            "DELETE FROM ch2_authors WHERE name LIKE 'copy_demo_author_%' RETURNING id"
        )
        print(f"  COPY 加载并清理 {cur.rowcount} 行")
    conn.commit()


# ------------------------------------------------------------------ #
#  Demo 6：动态构造表名 —— sql.Identifier 防注入
# ------------------------------------------------------------------ #
def demo_dynamic_identifier(conn: psycopg.Connection) -> None:
    section("Demo 6 · 动态拼接表名（用 psycopg.sql 模块，不要用 f-string！）")

    table = "ch2_books"
    with conn.cursor() as cur:
        query = sql.SQL("SELECT count(*) FROM {tbl}").format(
            tbl=sql.Identifier(table)
        )
        cur.execute(query)
        print(f"  表 {table} 共 {cur.fetchone()[0]} 行")


# ------------------------------------------------------------------ #
#  Demo 7：服务端游标 —— 流式遍历大结果集
# ------------------------------------------------------------------ #
def demo_server_cursor(conn: psycopg.Connection) -> None:
    section("Demo 7 · 服务端游标流式读 ch2_books 表（避免一次性加载到内存）")

    with conn.cursor(name="books_stream") as cur:
        cur.itersize = 5
        cur.execute("SELECT id, title FROM ch2_books ORDER BY id")
        for row in cur:
            print(f"  id={row[0]:<3}  title={row[1]}")


# ------------------------------------------------------------------ #
#  Main
# ------------------------------------------------------------------ #
def main() -> None:
    print(f"连接：{CONN_INFO}")
    try:
        conn = psycopg.connect(CONN_INFO)
    except psycopg.OperationalError as exc:
        sys.exit(f"连接失败：{exc}\n请先按第 2 章说明初始化 learn_pg 数据库与 init.sql。")

    try:
        with conn:
            demo_select_basic(conn)
            new_id = demo_insert_returning(conn)
            demo_transaction(conn, new_id)
            demo_executemany(conn)
            demo_copy(conn)
            demo_dynamic_identifier(conn)
            demo_server_cursor(conn)

            with conn.cursor() as cur:
                cur.execute(
                    "DELETE FROM ch2_books WHERE isbn = %s RETURNING id", ("9999999999999",)
                )
                if cur.rowcount:
                    print(f"\n  [清理] 删掉 demo 插入的书 id={cur.fetchone()[0]}")
            conn.commit()
    finally:
        conn.close()

    print("\n✅ All demos done.")


if __name__ == "__main__":
    main()