#!/usr/bin/env python3
"""Capture harness for DDIA Ch. 8 exercise: "Write skew defeats snapshot isolation".

Reproduces, against a REAL PostgreSQL 16, the on-call-doctors write skew from
DDIA 2e Ch. 8 (§"Write Skew and Phantoms"). Two doctors are on call; an
invariant requires at least one. Two concurrent transactions each check the
count (both see 2), each conclude it is safe to take *themselves* off call, and
each UPDATE a *different* row. The two versions of the story:

  1. REPEATABLE READ (Postgres snapshot isolation): both commit. Neither wrote
     the same row, so lost-update detection never fires. Final: 0 on call --
     the invariant is violated.

  2. SERIALIZABLE (Postgres SSI): the same interleaving; one commits, the other
     is aborted with SQLSTATE 40001, "could not serialize access due to
     read/write dependencies among transactions". Final: >=1 on call.

  3. FOR UPDATE at REPEATABLE READ: reading the counted rows FOR UPDATE takes
     row locks, so the second transaction blocks until the first commits, then
     re-reads and sees only 1 on call -- it must not leave. Final: >=1 on call.

Deterministic and idempotent: drops+recreates database `ex_ch08_writeskew`
every run and closes all connections before dropping. This is the maintainer's
verification path -- NOT the reader's path. The reader reproduces the same
thing by hand in two psql sessions (see write-skew-doctors.md, Steps).

Run:
    uv run --python 3.12 --with "psycopg[binary]" python3 code/write_skew_doctors.py
"""

import threading
import time

import psycopg
from psycopg import errors

ADMIN_DSN = "host=localhost port=5432 user=postgres dbname=study"
DB = "ex_ch08_writeskew"
DSN = f"host=localhost port=5432 user=postgres dbname={DB}"


def rule():
    print("=" * 70)


def reset_database():
    """Drop and recreate ex_ch08_writeskew via an autocommit admin connection."""
    with psycopg.connect(ADMIN_DSN, autocommit=True) as admin:
        admin.execute(f"DROP DATABASE IF EXISTS {DB} WITH (FORCE)")
        admin.execute(f"CREATE DATABASE {DB}")


def seed():
    """Create doctors(name, on_call, shift_id); both Alice and Bob on call, shift 1."""
    with psycopg.connect(DSN, autocommit=True) as conn:
        conn.execute(
            "CREATE TABLE doctors ("
            "  name text PRIMARY KEY,"
            "  on_call boolean NOT NULL,"
            "  shift_id int NOT NULL)"
        )
        conn.execute("INSERT INTO doctors VALUES ('Alice', true, 1)")
        conn.execute("INSERT INTO doctors VALUES ('Bob',   true, 1)")


def reseed():
    """Reset the table to the starting state between scenarios."""
    with psycopg.connect(DSN, autocommit=True) as conn:
        conn.execute("UPDATE doctors SET on_call = true WHERE shift_id = 1")


COUNT_SQL = "SELECT count(*) FROM doctors WHERE on_call = true AND shift_id = 1"


def final_on_call():
    with psycopg.connect(DSN, autocommit=True) as conn:
        return conn.execute(COUNT_SQL).fetchone()[0]


def scenario_repeatable_read():
    rule()
    print("SCENARIO 1  REPEATABLE READ -- the anomaly")
    rule()
    a = psycopg.connect(DSN)
    b = psycopg.connect(DSN)
    a.autocommit = False
    b.autocommit = False

    a.execute("BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ")
    b.execute("BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ")

    a_count = a.execute(COUNT_SQL).fetchone()[0]
    b_count = b.execute(COUNT_SQL).fetchone()[0]
    print("SESSION A:  BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ;")
    print(f"SESSION A:  {COUNT_SQL};")
    print(f"  -- count = {a_count}   (A sees >=2 on call, thinks it is safe to leave)")
    print("SESSION B:  BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ;")
    print(f"SESSION B:  {COUNT_SQL};")
    print(f"  -- count = {b_count}   (B sees >=2 on call, thinks it is safe to leave)")

    a.execute("UPDATE doctors SET on_call = false WHERE name = 'Alice'")
    b.execute("UPDATE doctors SET on_call = false WHERE name = 'Bob'")
    print("SESSION A:  UPDATE doctors SET on_call = false WHERE name = 'Alice';")
    print("SESSION B:  UPDATE doctors SET on_call = false WHERE name = 'Bob';")
    print("  -- different rows: no write-write conflict, no lost-update detection")

    a.commit()
    b.commit()
    print("SESSION A:  COMMIT;   -- succeeds")
    print("SESSION B:  COMMIT;   -- succeeds")

    a.close()
    b.close()

    remaining = final_on_call()
    print()
    print(f"  FINAL on_call count = {remaining}")
    if remaining == 0:
        print("  -- INVARIANT VIOLATED: zero doctors on call. Write skew.")
    else:
        print(f"  -- UNEXPECTED: expected 0, got {remaining}")
    return remaining


def scenario_serializable():
    reseed()
    rule()
    print("SCENARIO 2  SERIALIZABLE -- Postgres SSI catches it")
    rule()
    a = psycopg.connect(DSN)
    b = psycopg.connect(DSN)
    a.autocommit = False
    b.autocommit = False

    a.execute("BEGIN TRANSACTION ISOLATION LEVEL SERIALIZABLE")
    b.execute("BEGIN TRANSACTION ISOLATION LEVEL SERIALIZABLE")

    a_count = a.execute(COUNT_SQL).fetchone()[0]
    b_count = b.execute(COUNT_SQL).fetchone()[0]
    print("SESSION A:  BEGIN TRANSACTION ISOLATION LEVEL SERIALIZABLE;")
    print(f"SESSION A:  {COUNT_SQL};   -- count = {a_count}")
    print("SESSION B:  BEGIN TRANSACTION ISOLATION LEVEL SERIALIZABLE;")
    print(f"SESSION B:  {COUNT_SQL};   -- count = {b_count}")

    a.execute("UPDATE doctors SET on_call = false WHERE name = 'Alice'")
    b.execute("UPDATE doctors SET on_call = false WHERE name = 'Bob'")
    print("SESSION A:  UPDATE doctors SET on_call = false WHERE name = 'Alice';")
    print("SESSION B:  UPDATE doctors SET on_call = false WHERE name = 'Bob';")

    a.commit()
    print("SESSION A:  COMMIT;   -- succeeds")

    err_sqlstate = None
    err_text = None
    try:
        b.commit()
        print("SESSION B:  COMMIT;   -- UNEXPECTED: succeeded")
    except errors.SerializationFailure as e:
        err_sqlstate = e.sqlstate
        err_text = str(e).strip()
        print("SESSION B:  COMMIT;")
        print(f"  -- ERROR:  {err_sqlstate}")
        for line in err_text.splitlines():
            print(f"  -- {line}")
        b.rollback()

    a.close()
    b.close()

    remaining = final_on_call()
    print()
    print(f"  FINAL on_call count = {remaining}")
    if remaining >= 1 and err_sqlstate == "40001":
        print("  -- INVARIANT PRESERVED: SSI aborted one transaction (SQLSTATE 40001).")
    else:
        print(f"  -- UNEXPECTED: count={remaining}, sqlstate={err_sqlstate}")
    return remaining, err_sqlstate, err_text


def scenario_for_update():
    reseed()
    rule()
    print("SCENARIO 3  READ COMMITTED + SELECT ... FOR UPDATE -- lock the read")
    rule()
    # count(*) cannot take FOR UPDATE directly, so we lock the counted rows and
    # count what we locked. Session B runs in a real thread so its block is real.
    lock_sql = (
        "SELECT name FROM doctors WHERE on_call = true AND shift_id = 1 FOR UPDATE"
    )

    a = psycopg.connect(DSN)
    a.autocommit = False
    a.execute("BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED")
    a_rows = a.execute(lock_sql).fetchall()
    print("SESSION A:  BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;")
    print("SESSION A:  SELECT name ... WHERE on_call AND shift_id=1 FOR UPDATE;")
    print(f"  -- locked {len(a_rows)} rows: {sorted(r[0] for r in a_rows)}")

    b_result = {}

    def session_b():
        b = psycopg.connect(DSN)
        b.autocommit = False
        b.execute("BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED")
        t0 = time.monotonic()
        rows = b.execute(lock_sql).fetchall()  # blocks until A commits, then re-reads
        b_result["waited"] = time.monotonic() - t0
        b_result["rows"] = sorted(r[0] for r in rows)
        # B now sees only 1 on call, so it must NOT take itself off call.
        b.commit()
        b.close()

    print("SESSION B:  BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;")
    print("SESSION B:  SELECT name ... FOR UPDATE;   -- BLOCKS on A's row locks")
    tb = threading.Thread(target=session_b)
    tb.start()
    time.sleep(0.5)  # ensure B is genuinely waiting on the lock before A proceeds

    a.execute("UPDATE doctors SET on_call = false WHERE name = 'Alice'")
    a.commit()
    a.close()
    print("SESSION A:  UPDATE doctors SET on_call = false WHERE name = 'Alice';")
    print("SESSION A:  COMMIT;   -- releases the locks; B unblocks")

    tb.join()
    print(
        f"SESSION B:  (unblocked after ~{b_result['waited']:.1f}s)  "
        f"-- re-reads and now sees {len(b_result['rows'])} on call: {b_result['rows']}"
    )
    if len(b_result["rows"]) <= 1:
        print("  -- B sees only 1 on call: it must NOT take itself off call.")

    remaining = final_on_call()
    print()
    print(f"  FINAL on_call count = {remaining}")
    if remaining >= 1:
        print("  -- INVARIANT PRESERVED: FOR UPDATE serialized the decisions.")
    else:
        print(f"  -- UNEXPECTED: expected >=1, got {remaining}")
    return remaining


def main():
    reset_database()
    seed()

    rule()
    print("SETUP  doctors seeded, committed:")
    print("  ('Alice', on_call=true, shift_id=1)")
    print("  ('Bob',   on_call=true, shift_id=1)")
    print(f"  invariant: at least 1 on call for shift 1  (starts at 2)")

    rr = scenario_repeatable_read()
    ser, sqlstate, errtext = scenario_serializable()
    fu = scenario_for_update()

    rule()
    print("SUMMARY")
    print(f"  REPEATABLE READ  final on_call = {rr}   (0 = invariant violated)")
    print(f"  SERIALIZABLE     final on_call = {ser}   SQLSTATE = {sqlstate}")
    print(f"  FOR UPDATE       final on_call = {fu}   (>=1 = invariant preserved)")
    verdict = (
        "matches the aha"
        if rr == 0 and ser >= 1 and sqlstate == "40001" and fu >= 1
        else "DIVERGED from the aha -- rewrite the aha, keep the transcript"
    )
    print(f"  verdict          : {verdict}")
    rule()


if __name__ == "__main__":
    main()