# /// script
# requires-python = ">=3.10"
# dependencies = ["requests"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Per-record calls against batches: find the fixed cost.

A local HTTP server that does a trivial amount of work per record. 100,000
records sent as batches of 1 to 10,000, as a JSON array per call. Records
per second, the latency of a batch's first record, and the fit
time per call = F + c x batch size.
"""
import argparse, json, threading, time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import requests

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--records", type=int, default=100_000, help="Records sent per batch size.")
p.add_argument("--sizes", type=lambda s: [int(v) for v in s.split(",")], default=[1, 10, 100, 1000, 10000], help="Batch sizes.")
p.add_argument("--max-calls", type=int, default=5000, help="Cap on calls per size (batch size 1 sends this many records).")
p.add_argument("--port", type=int, default=8139, help="Port for the local server.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

class H(BaseHTTPRequestHandler):
    def do_POST(self):
        body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
        total = sum(r["qty"] * r["price"] for r in body)                        # trivial work per record
        out = json.dumps({"n": len(body), "total": total}).encode()
        self.send_response(200); self.send_header("Content-Length", str(len(out))); self.end_headers(); self.wfile.write(out)
    def log_message(self, *a): pass

srv = ThreadingHTTPServer(("127.0.0.1", args.port), H)
threading.Thread(target=srv.serve_forever, daemon=True).start()
s = requests.Session(); url = f"http://127.0.0.1:{args.port}/"

rows = []
for b in args.sizes:
    n = min(args.records, b * args.max_calls)
    batch = [{"id": i, "qty": 2, "price": 9.5} for i in range(b)]
    t0 = time.perf_counter()
    for _ in range(n // b): s.post(url, json=batch)
    dt = time.perf_counter() - t0
    calls = n // b
    rows.append({"batch": b, "records_per_s": n / dt, "per_call_ms": dt / calls * 1000, "first_record_ms": dt / calls * 1000})
    print(f"batch {b:>6,}: {n / dt:>10,.0f} records/s   {dt / calls * 1000:7.2f} ms per call (the first record waits for the whole call)")
srv.shutdown()
# fit F + c*B from the smallest and largest batch
a, z = rows[0], rows[-1]
c = (z["per_call_ms"] - a["per_call_ms"]) / (z["batch"] - a["batch"]); F = a["per_call_ms"] - c * a["batch"]
print(f"fit: F = {F:.3f} ms per call, c = {c * 1000:.2f} µs per record; F equals the per-record cost at a batch of {F / c:,.0f}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"batches": rows, "summary": [{"F_ms": F, "c_us": c * 1000, "crossover": F / c}]}}, f)