# /// script
# requires-python = ">=3.10"
# dependencies = ["flask", "requests"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Duplicates against an idempotency key.

A payment service with an in-memory ledger; a client sends 1,000 charges and,
with probability 0.2, the reply is lost (the server still processed the
request) and the client retries. Version 1: a new random key per attempt.
Version 2: the key is the intent, and the server replays the stored outcome.
Version 3: the server suppresses the repeat but answers with an error.
"""
import argparse, json, logging, random, threading, time, uuid
import requests
from flask import Flask, request, jsonify
from werkzeug.serving import make_server

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--charges", type=int, default=1000, help="Charges.")
p.add_argument("--drop", type=float, default=0.2, help="Probability that a reply is lost.")
p.add_argument("--port", type=int, default=8141, help="Port for the local server.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
random.seed(args.seed)
logging.getLogger("werkzeug").setLevel(logging.ERROR)
app = Flask(__name__)
state = {"ledger": [], "seen": {}, "mode": 2}
lock = threading.Lock()

@app.post("/charge")
def charge():
    key = request.headers["Idempotency-Key"]; amount = request.json["amount"]
    with lock:
        if key in state["seen"]:
            if state["mode"] == 3: return jsonify({"status": "error", "reason": "duplicate"}), 409
            return jsonify(state["seen"][key])                  # replay the stored outcome
        state["ledger"].append(amount)
        state["seen"][key] = {"status": "ok", "n": len(state["ledger"])}
        out = state["seen"][key]
    if random.random() < args.drop: time.sleep(0.6)             # the reply is lost: it arrives after the client gives up
    return jsonify(out)

srv = make_server("127.0.0.1", args.port, app, threaded=True); threading.Thread(target=srv.serve_forever, daemon=True).start()
url = f"http://127.0.0.1:{args.port}/charge"

def run(mode):
    with lock: state["ledger"].clear(); state["seen"].clear(); state["mode"] = 2 if mode == 1 else mode
    s, errors = requests.Session(), 0
    for i in range(args.charges):
        intent = f"result-{i}"
        while True:
            key = str(uuid.uuid4()) if mode == 1 else intent
            try:
                r = s.post(url, json={"amount": 10}, headers={"Idempotency-Key": key}, timeout=0.3)
                if r.status_code != 200: errors += 1
                break
            except requests.exceptions.Timeout: pass
    with lock: rows = len(state["ledger"])
    return {"version": {1: "1: new key per attempt", 2: "2: key = intent, replay outcome", 3: "3: suppress, answer with an error"}[mode],
            "ledger_rows": rows, "extra_rows": rows - args.charges, "client_errors": errors}

rows = [run(m) for m in (1, 2, 3)]
for r in rows: print(f"{r['version']:<34} ledger rows {r['ledger_rows']:>5} ({r['extra_rows']:+d})   errors returned to callers {r['client_errors']}")
srv.shutdown()

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"versions": rows}}, f)