"""
02_isolation_levels.py —— 不同隔离级别下的并发现象

依赖：
    pip install "psycopg[binary]>=3.1"

用法：
    1. 初始化：psql -h 127.0.0.1 -U postgres -d learn_pg -f ../init.sql
    2. python 02_isolation_levels.py [rc | rr | sz]
       - rc：READ COMMITTED 下复现「不可重复读」
       - rr：REPEATABLE READ 下证明「不可重复读消失」
       - sz：SERIALIZABLE 下复现「写偏序被回滚」

通过两个 psycopg 连接模拟两个并发会话 A、B。
"""
from __future__ import annotations

import sys
import time

import psycopg
from psycopg import IsolationLevel
from psycopg.errors import SerializationFailure

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

ISO_MAP = {
    "rc": IsolationLevel.READ_COMMITTED,
    "rr": IsolationLevel.REPEATABLE_READ,
    "sz": IsolationLevel.SERIALIZABLE,
}


def reset_data() -> None:
    with psycopg.connect(DSN, autocommit=True) as conn:
        with conn.cursor() as cur:
            cur.execute("UPDATE ch7_products SET price=7999, stock=10 WHERE id=1")
            cur.execute("UPDATE ch7_products SET price=4999, stock=20 WHERE id=2")
            cur.execute("UPDATE ch7_doctors SET on_duty=TRUE")


def open_conn(level: IsolationLevel) -> psycopg.Connection:
    conn = psycopg.connect(DSN, autocommit=False)
    conn.isolation_level = level
    return conn


def demo_non_repeatable_read(level: IsolationLevel) -> None:
    """同一事务两次 SELECT，看是否一致。"""
    label = level.name
    print(f"\n=== [{label}] 不可重复读测试 ===")
    reset_data()
    a = open_conn(level)
    b = open_conn(level)

    with a.cursor() as ca, b.cursor() as cb:
        ca.execute("SELECT price FROM ch7_products WHERE id=1")
        p1 = ca.fetchone()[0]
        print(f"  [A] 第 1 次 SELECT price = {p1}")

        cb.execute("UPDATE ch7_products SET price = 8888 WHERE id=1")
        b.commit()
        print(f"  [B] UPDATE price -> 8888 并 COMMIT")

        ca.execute("SELECT price FROM ch7_products WHERE id=1")
        p2 = ca.fetchone()[0]
        print(f"  [A] 第 2 次 SELECT price = {p2}")

        if p1 == p2:
            print(f"  ✓ 两次读一致 -> {label} 防住了不可重复读")
        else:
            print(f"  ✗ 两次读不一致 -> {label} 没能防住不可重复读")
        a.commit()

    a.close()
    b.close()


def demo_write_skew(level: IsolationLevel) -> None:
    """医生值班表的写偏序场景。"""
    label = level.name
    print(f"\n=== [{label}] 写偏序（医生值班）测试 ===")
    reset_data()
    a = open_conn(level)
    b = open_conn(level)

    try:
        with a.cursor() as ca, b.cursor() as cb:
            ca.execute("SELECT count(*) FROM ch7_doctors WHERE on_duty")
            n_a = ca.fetchone()[0]
            cb.execute("SELECT count(*) FROM ch7_doctors WHERE on_duty")
            n_b = cb.fetchone()[0]
            print(f"  [A] 在班医生 = {n_a}，决定让自己请假")
            print(f"  [B] 在班医生 = {n_b}，决定让自己请假")

            ca.execute("UPDATE ch7_doctors SET on_duty=FALSE WHERE id=1")
            cb.execute("UPDATE ch7_doctors SET on_duty=FALSE WHERE id=2")

            try:
                a.commit()
                print(f"  [A] COMMIT 成功")
            except SerializationFailure as e:
                print(f"  [A] COMMIT 失败：{e}")

            try:
                b.commit()
                print(f"  [B] COMMIT 成功")
            except SerializationFailure as e:
                print(f"  [B] COMMIT 失败（被 SSI 保护）：{e.diag.message_primary}")
    finally:
        a.close()
        b.close()

    with psycopg.connect(DSN) as conn:
        cur = conn.execute("SELECT count(*) FROM ch7_doctors WHERE on_duty")
        n = cur.fetchone()[0]
        print(f"  最终在班医生 = {n}（业务规则：必须 ≥ 1）")


def main() -> None:
    arg = sys.argv[1] if len(sys.argv) > 1 else "all"

    if arg in ("rc", "all"):
        demo_non_repeatable_read(IsolationLevel.READ_COMMITTED)
    if arg in ("rr", "all"):
        demo_non_repeatable_read(IsolationLevel.REPEATABLE_READ)
        demo_write_skew(IsolationLevel.REPEATABLE_READ)
    if arg in ("sz", "all"):
        demo_write_skew(IsolationLevel.SERIALIZABLE)


if __name__ == "__main__":
    main()