"""
Ch4 配套代码 3 / 3 —— SERIAL vs IDENTITY 完整对比

依赖：先 init.sql 建好 ch4_serial_demo / ch4_identity_always / ch4_identity_default

演示：
  1. SERIAL 允许直接塞 id（埋下后续主键冲突的雷）
  2. IDENTITY ALWAYS 拒绝直接塞 id
  3. OVERRIDING SYSTEM VALUE 强制覆盖
  4. 序列与 IDENTITY 的 nextval / setval / RESTART
  5. 事务回滚导致序列跳号（这是 PG 的设计取舍）
"""

import psycopg
from psycopg import errors as pe


DSN = "host=127.0.0.1 port=5432 dbname=learn_pg user=postgres"


def section(title: str) -> None:
    print("\n" + "=" * 64)
    print(title)
    print("=" * 64)


def demo_serial_can_force(conn: psycopg.Connection) -> None:
    section("Demo 1: SERIAL 列允许手工指定 id —— 危险！")

    with conn.cursor() as cur:
        cur.execute("TRUNCATE ch4_serial_demo RESTART IDENTITY")

        # 正常插入，让序列推进
        cur.execute("INSERT INTO ch4_serial_demo (name) VALUES ('row1') RETURNING id")
        print(f"  自动 id = {cur.fetchone()[0]}  (序列给的 1)")

        # 手工塞一个未来的 id —— SERIAL 不会拦
        cur.execute("INSERT INTO ch4_serial_demo (id, name) VALUES (10, 'manual') RETURNING id")
        print(f"  手工 id = {cur.fetchone()[0]}  (硬塞，序列没动！)")

        # 再正常插入，序列会从 2 开始 —— 但其实表里 id=2 没人用，没问题
        # 然后继续到 id=10 时才会爆炸
        cur.execute("INSERT INTO ch4_serial_demo (name) VALUES ('row3') RETURNING id")
        print(f"  自动 id = {cur.fetchone()[0]}  (序列给的 2，与手工塞的 10 暂时不冲突)")

        # 把序列设置为 9，模拟「快撞了」
        cur.execute("SELECT setval(pg_get_serial_sequence('ch4_serial_demo', 'id'), 9)")

        cur.execute("INSERT INTO ch4_serial_demo (name) VALUES ('row10') RETURNING id")
        print(f"  自动 id = {cur.fetchone()[0]}  (序列推进到 10)")

        try:
            cur.execute("INSERT INTO ch4_serial_demo (name) VALUES ('boom') RETURNING id")
            print(f"  ⚠️ 未爆炸: id = {cur.fetchone()[0]}")
        except pe.UniqueViolation as e:
            conn.rollback()
            print(f"  💥 主键冲突! {str(e).splitlines()[0]}")
            print("     -> SERIAL 让手工塞值与序列脱节，运维事故的常见根源")


def demo_identity_always_rejects(conn: psycopg.Connection) -> None:
    section("Demo 2: IDENTITY ALWAYS 拒绝手工塞 id —— 安全！")

    with conn.cursor() as cur:
        cur.execute("TRUNCATE ch4_identity_always RESTART IDENTITY")

        cur.execute("INSERT INTO ch4_identity_always (name) VALUES ('row1') RETURNING id")
        print(f"  自动 id = {cur.fetchone()[0]}")

        try:
            cur.execute("INSERT INTO ch4_identity_always (id, name) VALUES (100, 'manual')")
            print("  ⚠️ 未预期通过！")
        except pe.GeneratedAlways as e:
            conn.rollback()
            print(f"  ✅ 拒绝手工塞: {str(e).splitlines()[0]}")

        # OVERRIDING SYSTEM VALUE 强制覆盖（数据导入场景）
        cur.execute(
            "INSERT INTO ch4_identity_always (id, name) "
            "OVERRIDING SYSTEM VALUE VALUES (1000, 'imported') RETURNING id"
        )
        print(f"  ✅ OVERRIDING SYSTEM VALUE 强制塞: id = {cur.fetchone()[0]}")
        conn.commit()


def demo_identity_default(conn: psycopg.Connection) -> None:
    section("Demo 3: IDENTITY BY DEFAULT —— 行为接近 SERIAL")

    with conn.cursor() as cur:
        cur.execute("TRUNCATE ch4_identity_default RESTART IDENTITY")

        cur.execute("INSERT INTO ch4_identity_default (name) VALUES ('a') RETURNING id")
        print(f"  自动 id = {cur.fetchone()[0]}")
        cur.execute("INSERT INTO ch4_identity_default (id, name) VALUES (50, 'force') RETURNING id")
        print(f"  手工 id = {cur.fetchone()[0]}  (BY DEFAULT 允许)")
        conn.commit()


def demo_sequence_ops(conn: psycopg.Connection) -> None:
    section("Demo 4: 序列函数 nextval / currval / setval / lastval")

    with conn.cursor() as cur:
        cur.execute("SELECT nextval('ch4_order_no_seq')")
        a = cur.fetchone()[0]
        cur.execute("SELECT nextval('ch4_order_no_seq')")
        b = cur.fetchone()[0]
        cur.execute("SELECT currval('ch4_order_no_seq')")
        c = cur.fetchone()[0]
        cur.execute("SELECT lastval()")
        d = cur.fetchone()[0]
        print(f"  连续两次 nextval: {a}, {b}")
        print(f"  currval (本会话最近一次该序列): {c}")
        print(f"  lastval (本会话最近一次任意序列): {d}")

        cur.execute("SELECT setval('ch4_order_no_seq', 99000)")
        cur.execute("SELECT nextval('ch4_order_no_seq')")
        print(f"  setval 到 99000 后再 nextval: {cur.fetchone()[0]}")


def demo_rollback_skip(conn: psycopg.Connection) -> None:
    section("Demo 5: 事务回滚导致序列跳号（PG 设计取舍）")

    with conn.cursor() as cur:
        cur.execute("TRUNCATE ch4_identity_default RESTART IDENTITY")

        # 事务 1：成功
        cur.execute("INSERT INTO ch4_identity_default (name) VALUES ('keep') RETURNING id")
        print(f"  事务 1 成功插入 id = {cur.fetchone()[0]}")
        conn.commit()

        # 事务 2：插入再回滚 —— 但序列不会回退！
        cur.execute("INSERT INTO ch4_identity_default (name) VALUES ('rollback') RETURNING id")
        print(f"  事务 2 插入 id = {cur.fetchone()[0]}")
        conn.rollback()
        print("  事务 2 ROLLBACK")

        # 事务 3：再插入，看 id 是否跳号
        cur.execute("INSERT INTO ch4_identity_default (name) VALUES ('after') RETURNING id")
        print(f"  事务 3 插入 id = {cur.fetchone()[0]}    ← 跳过了被回滚的 id")
        conn.commit()

        print()
        print("  💡 这是 PG 的有意设计：nextval 不受事务保护，避免并发回滚")
        print("     互相阻塞。代价是 id 会跳号。如果需要严格连续号请用其他方案")
        print("     （如表锁 + max(id)+1，或订单号生成器）。")


def main() -> None:
    with psycopg.connect(DSN, autocommit=False) as conn:
        demo_serial_can_force(conn)
        demo_identity_always_rejects(conn)
        demo_identity_default(conn)
        demo_sequence_ops(conn)
        demo_rollback_skip(conn)


if __name__ == "__main__":
    try:
        main()
    except psycopg.OperationalError as e:
        print(f"❌ 连接 PostgreSQL 失败: {e}")