#!/usr/bin/env python3
"""Back-of-envelope calculators for system design rounds.

The point is not the tool. The point is that you do this arithmetic OUT LOUD in
the round, and that you have done it often enough that the numbers come without
a calculator. Use this to check yourself, and to build the reflex.

    python3 envelope.py qps       --rps 50000 --ms 8
    python3 envelope.py queue     --rho 0.8
    python3 envelope.py storage   --rows 1e10 --bytes 512 --replicas 3
    python3 envelope.py network   --rps 50000 --bytes 4096
    python3 envelope.py retry     --rps 10000 --fail 0.3 --attempts 3
    python3 envelope.py latencies
    python3 envelope.py kv        --keys 1e9 --value-bytes 200
"""

from __future__ import annotations

import argparse


def human(n: float, unit: str = "") -> str:
    for suffix, scale in (("T", 1e12), ("G", 1e9), ("M", 1e6), ("K", 1e3)):
        if abs(n) >= scale:
            return f"{n / scale:.2f} {suffix}{unit}"
    return f"{n:.2f} {unit}".strip()


def bytes_human(n: float) -> str:
    for suffix, scale in (("PiB", 2**50), ("TiB", 2**40), ("GiB", 2**30),
                          ("MiB", 2**20), ("KiB", 2**10)):
        if abs(n) >= scale:
            return f"{n / scale:.2f} {suffix}"
    return f"{n:.0f} B"


def rule(title: str) -> None:
    print(f"\n{title}\n{'-' * max(len(title), 40)}")


# ---------------------------------------------------------------------------


def cmd_qps(args) -> None:
    """Little's law: concurrency = arrival rate x service time."""
    rps, seconds = args.rps, args.ms / 1000.0
    concurrency = rps * seconds

    rule("Little's law — L = λW")
    print(f"  arrival rate λ   {human(rps, 'req/s')}")
    print(f"  service time W   {args.ms:.1f} ms")
    print(f"  concurrency L    {concurrency:,.0f} requests in flight")

    rule("Capacity")
    per_worker = args.per_worker
    workers = concurrency / per_worker
    print(f"  at {per_worker} concurrent per worker -> {workers:,.1f} workers "
          f"(bare minimum)")
    print(f"  + headroom to {args.target_util:.0%} utilization -> "
          f"{workers / args.target_util:,.1f} steady-state")
    print(f"  + survive losing 1 of 3 AZs (x1.5)   -> "
          f"{workers / args.target_util * 1.5:,.0f} provisioned")

    rule("If it is CPU-bound instead")
    cpu_ms = args.cpu_ms if args.cpu_ms else args.ms
    cores = rps * (cpu_ms / 1000.0)
    print(f"  at {cpu_ms:.1f} ms CPU per request -> {cores:,.1f} cores busy")
    print(f"  at {args.target_util:.0%} target utilization -> "
          f"{cores / args.target_util:,.0f} cores")
    print(f"  on {args.cores_per_box}-core boxes    -> "
          f"{cores / args.target_util / args.cores_per_box:,.1f} boxes")

    print("\n  Say this out loud in the round. Any arithmetic beats none.")


def cmd_queue(args) -> None:
    """The utilization/latency knee for an M/M/1 queue."""
    rule("Why 80% utilization is not '80% as bad as 100%'")
    print("  M/M/1: waiting time multiplier = ρ / (1 - ρ)")
    print("  Total response time = service time x 1/(1-ρ)\n")
    print(f"  {'ρ':>6}  {'wait / service':>16}  {'total / service':>17}")
    for rho in (0.5, 0.7, 0.8, 0.9, 0.95, 0.99, 0.999):
        wait = rho / (1 - rho)
        print(f"  {rho:>6.3f}  {wait:>16.1f}x  {1 / (1 - rho):>16.1f}x")

    rho = args.rho
    rule(f"At your ρ = {rho}")
    print(f"  queue wait is {rho / (1 - rho):.2f}x the service time")
    print(f"  response time is {1 / (1 - rho):.2f}x the service time")
    delta = 0.05
    new = min(rho + delta, 0.999)
    print(f"  a {delta:.0%} traffic increase (ρ -> {new:.2f}) makes it "
          f"{1 / (1 - new):.2f}x — a {(1 / (1 - new)) / (1 / (1 - rho)):.2f}x jump")
    print("""
  This is THE argument for load shedding, and it is the reason a system at 85%
  looks fine on a dashboard and falls over at 92%. Latency is not linear in
  utilization; it is hyperbolic. Every capacity plan must state its target ρ,
  and anything above ~0.8 needs admission control rather than more patience.

  Caveat to state if pressed: M/M/1 assumes Poisson arrivals and exponential
  service times. Real traffic is burstier, so the real knee arrives EARLIER
  than this table suggests, not later.""")


def cmd_storage(args) -> None:
    rows, per_row, replicas = args.rows, args.bytes, args.replicas
    raw = rows * per_row

    rule("Storage")
    print(f"  {human(rows)} rows x {per_row} B = {bytes_human(raw)} logical")
    print(f"  x {replicas} replicas              = {bytes_human(raw * replicas)}")
    print(f"  x {args.overhead} index/WAL/compaction overhead = "
          f"{bytes_human(raw * replicas * args.overhead)}")

    if args.growth_per_day:
        daily = args.growth_per_day * per_row * replicas * args.overhead
        print(f"\n  growth {human(args.growth_per_day)} rows/day = "
              f"{bytes_human(daily)}/day")
        print(f"  one year                     = {bytes_human(daily * 365)}")

    per_node = args.node_tb * 2**40
    print(f"\n  at {args.node_tb} TiB usable per node -> "
          f"{raw * replicas * args.overhead / per_node:,.1f} nodes")
    print("\n  Always say which of these you mean: logical, replicated, or on-disk.")


def cmd_network(args) -> None:
    rps, size = args.rps, args.bytes
    bps = rps * size * 8

    rule("Network")
    print(f"  {human(rps, 'req/s')} x {bytes_human(size)} = "
          f"{bytes_human(rps * size)}/s = {human(bps, 'bit/s')}")
    print(f"  on a 10 Gbit/s NIC -> {bps / 10e9:.1%} of one NIC")
    print(f"  on a 25 Gbit/s NIC -> {bps / 25e9:.1%} of one NIC")
    print(f"  per day            -> {bytes_human(rps * size * 86400)}")
    print(f"  egress at ${args.egress_per_gb}/GiB -> "
          f"${rps * size * 86400 * 30 / 2**30 * args.egress_per_gb:,.0f}/month")
    print("\n  Cross-AZ and egress costs are where designs get expensive quietly.")


def cmd_retry(args) -> None:
    rps, fail, attempts = args.rps, args.fail, args.attempts

    rule("Retry amplification")
    total = rps * sum(fail**i for i in range(attempts))
    print(f"  base load          {human(rps, 'req/s')}")
    print(f"  failure rate       {fail:.0%}")
    print(f"  up to {attempts} attempts")
    print(f"  offered load       {human(total, 'req/s')}  "
          f"({total / rps:.2f}x amplification)")

    print("\n  Now the same thing during an incident:")
    for f in (0.5, 0.8, 0.95, 1.0):
        amp = sum(f**i for i in range(attempts))
        print(f"    failure {f:>4.0%} -> {amp:.2f}x offered load "
              f"= {human(rps * amp, 'req/s')}")
    print(f"""
  A dependency that degrades to 95% failure sees {sum(0.95**i for i in range(attempts)):.1f}x
  its normal traffic from retries alone — at the exact moment it is least able
  to serve it. That is a retry storm, and it converts a partial outage into a
  total one.

  The fixes, in order of importance:
    1. RETRY BUDGET — cap retries at a fraction (say 10%) of base traffic, so
       amplification is bounded no matter how bad it gets.
    2. CIRCUIT BREAKER — stop trying entirely, then probe with a half-open.
    3. JITTER — without it, every client retries in lockstep and you get a
       synchronized thundering herd the moment the dependency recovers.
  Jitter alone is not enough. Budget first.""")


def cmd_kv(args) -> None:
    keys, value = args.keys, args.value_bytes
    key_bytes = args.key_bytes
    entry = key_bytes + value + args.per_entry_overhead

    rule("Key-value sizing")
    print(f"  {human(keys)} keys x ({key_bytes} + {value} + "
          f"{args.per_entry_overhead}) B = {bytes_human(keys * entry)}")
    print(f"  in-memory index only ({key_bytes} + 8 B ptr) = "
          f"{bytes_human(keys * (key_bytes + 8))}")
    print(f"\n  Full dataset needs {bytes_human(keys * entry)}; the index alone needs")
    print(f"  {bytes_human(keys * (key_bytes + 8))}. Compare both against your node memory.")
    print(f"  If the INDEX does not fit in RAM, you are designing an on-disk structure")
    print(f"  (LSM or B-tree), not a hash map — and that changes the entire read path,")
    print(f"  because every lookup now costs a disk seek instead of a pointer chase.")


def cmd_latencies(args) -> None:
    rule("Numbers to have memorized")
    data = [
        ("L1 cache reference", "1 ns"),
        ("Branch mispredict", "3 ns"),
        ("L2 cache reference", "4 ns"),
        ("Mutex lock/unlock", "17 ns"),
        ("Main memory reference", "100 ns"),
        ("Compress 1 KB (snappy)", "2 µs"),
        ("Read 1 MB sequentially from memory", "3 µs"),
        ("SSD random read", "16–100 µs"),
        ("Read 1 MB sequentially from SSD", "49 µs"),
        ("Round trip within same datacenter", "500 µs"),
        ("Read 1 MB sequentially from disk", "825 µs"),
        ("Disk seek (spinning)", "10 ms"),
        ("Round trip US cross-country", "40–70 ms"),
        ("Round trip US ↔ Europe", "80–150 ms"),
    ]
    for label, value in data:
        print(f"  {label:<40} {value:>12}")
    print("""
  Two derived rules worth having cold:
    - Memory is ~100x faster than SSD; SSD is ~100x faster than a disk seek.
    - Anything crossing a datacenter boundary costs at least 0.5 ms, so a
      design with five sequential cross-service hops has a 2.5 ms floor before
      it does any work at all. That is why fan-out beats chaining.

  Source: Jeff Dean's 'Latency Numbers Every Programmer Should Know', adjusted
  for modern SSDs. Orders of magnitude, not benchmarks — quote them as such.""")


# ---------------------------------------------------------------------------


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = parser.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("qps", help="Little's law, concurrency, cores")
    p.add_argument("--rps", type=float, required=True)
    p.add_argument("--ms", type=float, required=True, help="service latency, ms")
    p.add_argument("--cpu-ms", type=float, default=None, help="CPU time per request, ms")
    p.add_argument("--per-worker", type=int, default=200)
    p.add_argument("--cores-per-box", type=int, default=16)
    p.add_argument("--target-util", type=float, default=0.6)
    p.set_defaults(func=cmd_qps)

    p = sub.add_parser("queue", help="the utilization/latency knee")
    p.add_argument("--rho", type=float, default=0.8)
    p.set_defaults(func=cmd_queue)

    p = sub.add_parser("storage", help="bytes, replicated and amplified")
    p.add_argument("--rows", type=float, required=True)
    p.add_argument("--bytes", type=float, required=True)
    p.add_argument("--replicas", type=int, default=3)
    p.add_argument("--overhead", type=float, default=1.5)
    p.add_argument("--growth-per-day", type=float, default=0)
    p.add_argument("--node-tb", type=float, default=8)
    p.set_defaults(func=cmd_storage)

    p = sub.add_parser("network", help="bandwidth and egress cost")
    p.add_argument("--rps", type=float, required=True)
    p.add_argument("--bytes", type=float, required=True)
    p.add_argument("--egress-per-gb", type=float, default=0.09)
    p.set_defaults(func=cmd_network)

    p = sub.add_parser("retry", help="retry amplification and storms")
    p.add_argument("--rps", type=float, required=True)
    p.add_argument("--fail", type=float, default=0.1)
    p.add_argument("--attempts", type=int, default=3)
    p.set_defaults(func=cmd_retry)

    p = sub.add_parser("kv", help="key-value sizing")
    p.add_argument("--keys", type=float, required=True)
    p.add_argument("--value-bytes", type=float, default=200)
    p.add_argument("--key-bytes", type=float, default=32)
    p.add_argument("--per-entry-overhead", type=float, default=48)
    p.set_defaults(func=cmd_kv)

    p = sub.add_parser("latencies", help="the memorized table")
    p.set_defaults(func=cmd_latencies)

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()