#!/usr/bin/env python3
"""Capture harness for the DDIA Ch. 8 read-skew exercise.

Reproduces read skew (a non-repeatable read) on a REAL PostgreSQL 16 server,
under the default READ COMMITTED isolation level, using the book's two-accounts
example (Aaliyah, $500 + $500). Then re-runs the identical interleaving under
REPEATABLE READ and shows the anomaly is gone.

This is the CAPTURE / VERIFY path, not the reader path. Readers of the exercise
type two `psql` sessions by hand (Session A / Session B); this script drives the
same two sessions programmatically so the printed transcript is real, verbatim
output that the write-up quotes.

It is deterministic and idempotent: it drops and recreates its own database
`ex_ch08_readskew` on the shared, already-running server, seeds it, runs both
scenarios, and closes every connection before dropping.

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

Substrate (already running; this script does NOT manage containers):
    PostgreSQL 16 container `ddia-ch08-pg` on localhost:5432, user=postgres,
    trust auth, maintenance db `study`.
"""

import psycopg

ADMIN_DSN = "host=localhost port=5432 user=postgres dbname=study"
DBNAME = "ex_ch08_readskew"
EX_DSN = f"host=localhost port=5432 user=postgres dbname={DBNAME}"


def recreate_database() -> None:
    """Drop and recreate our own database via an autocommit admin connection."""
    with psycopg.connect(ADMIN_DSN, autocommit=True) as admin:
        admin.execute(f"DROP DATABASE IF EXISTS {DBNAME} WITH (FORCE)")
        admin.execute(f"CREATE DATABASE {DBNAME}")


def seed() -> None:
    """Create the accounts table and seed two accounts at $500 each."""
    with psycopg.connect(EX_DSN, autocommit=True) as conn:
        conn.execute("DROP TABLE IF EXISTS accounts")
        conn.execute(
            "CREATE TABLE accounts (id int PRIMARY KEY, name text, balance int)"
        )
        conn.execute(
            "INSERT INTO accounts VALUES (1, 'checking', 500), (2, 'savings', 500)"
        )


def reset_balances() -> None:
    """Return both accounts to $500 between scenarios."""
    with psycopg.connect(EX_DSN, autocommit=True) as conn:
        conn.execute("UPDATE accounts SET balance = 500")


def balance(cur, acct_id: int) -> int:
    cur.execute("SELECT balance FROM accounts WHERE id = %s", (acct_id,))
    return cur.fetchone()[0]


def transfer_100_savings_to_checking() -> None:
    """Session B: move $100 from savings (2) to checking (1), atomically committed.

    A fresh connection, its own transaction, committed in full before it returns
    -- exactly what the reader does by hand in Session B.
    """
    with psycopg.connect(EX_DSN, autocommit=False) as sb:
        with sb.cursor() as cur:
            cur.execute("UPDATE accounts SET balance = balance - 100 WHERE id = 2")
            cur.execute("UPDATE accounts SET balance = balance + 100 WHERE id = 1")
        sb.commit()


def run_scenario(level: str) -> int:
    """Run the read-skew interleaving at the given isolation level for Session A.

    Session A reads checking, THEN Session B transfers 100 savings->checking and
    commits, THEN Session A reads savings. Returns Session A's summed total.
    """
    label = level
    print(f"=== Session A isolation level: {label} ===")
    with psycopg.connect(EX_DSN, autocommit=False) as sa:
        sa.execute(f"SET TRANSACTION ISOLATION LEVEL {level}")
        with sa.cursor() as cur:
            # Open Session A's transaction with its first read.
            checking = balance(cur, 1)
            print(f"Session A  read 1: SELECT balance WHERE id=1 (checking)  ->  {checking}")

            # Session B transfers $100 savings->checking and commits, in the gap
            # between Session A's two reads.
            transfer_100_savings_to_checking()
            print("Session B  committed: 100 moved savings(2) -> checking(1)")

            # Session A's second read, still in the same transaction.
            savings = balance(cur, 2)
            print(f"Session A  read 2: SELECT balance WHERE id=2 (savings)   ->  {savings}")

            total = checking + savings
            print(f"Session A  total seen: {checking} + {savings} = {total}")
        sa.rollback()
    print()
    return total


def main() -> None:
    recreate_database()
    seed()

    print("Two accounts, $500 each; true total is always $1000.\n")

    # Scenario 1: the default. Read skew appears.
    reset_balances()
    rc_total = run_scenario("READ COMMITTED")

    # Scenario 2: the fix. Snapshot fixed at first query; no skew.
    reset_balances()
    rr_total = run_scenario("REPEATABLE READ")

    print("=== summary ===")
    print(f"READ COMMITTED  Session A summed to  {rc_total}"
          f"  ({'read skew: $100 vanished' if rc_total != 1000 else 'consistent'})")
    print(f"REPEATABLE READ Session A summed to  {rr_total}"
          f"  ({'consistent' if rr_total == 1000 else 'read skew: $100 vanished'})")


if __name__ == "__main__":
    main()