#!/usr/bin/env bash
# Real demos for P004 "Every System Design Concept You Need, in One Diagram".
# Runs on render01 (Docker): DNS, a load balancer over 3 servers, cache vs database,
# an index, primary/replica replication with failover, a message queue, a CDN header
# check, hash sharding and a load test. Every command + its real output is saved to
# demo_output.json; terminal scenes show only this captured output.
# Every number in video.json must match a value here. Re-run = new numbers = update the script.
set -euo pipefail
cd "$(dirname "$0")"
P=ce-p004
cleanup() {
  docker rm -f lb server-1 server-2 server-3 db redis primary replica py >/dev/null 2>&1 || true
  docker network rm $P >/dev/null 2>&1 || true
}
cleanup
trap cleanup EXIT
docker network create $P >/dev/null

python3 - "$P" <<'PY'
import json, re, subprocess, sys, time
P = sys.argv[1]
steps = {}

def sh(cmd, check=False):
    r = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    if check and r.returncode:
        raise SystemExit(f"FAILED: {cmd}\n{r.stdout}\n{r.stderr}")
    return (r.stdout + r.stderr).rstrip()

def step(key, show, real=None, out=None):
    """Run `real` (default: `show`) and record the command as the viewer sees it."""
    steps[key] = {"cmd": show, "out": out if out is not None else sh(real or show)}
    print(f"== {key}\n$ {show}\n{steps[key]['out']}\n", flush=True)

def wait(cmd, tries=60):
    for _ in range(tries):
        if subprocess.run(cmd, shell=True, capture_output=True).returncode == 0:
            return
        time.sleep(1)
    raise SystemExit(f"timed out waiting for: {cmd}")

# ---------------------------------------------------------------- 1. DNS
step("dns", "dig @1.1.1.1 wikipedia.org", "dig @1.1.1.1 +noall +answer +stats wikipedia.org | grep -E 'IN\\s+A|Query time'")

# ---------------------------------------------------------------- 2. load balancer
for i in (1, 2, 3):
    sh(f"docker run -d --name server-{i} --hostname server-{i} --network {P} traefik/whoami:v1.10", check=True)
conf = f"""events {{}}
http {{
  upstream app {{ server server-1; server server-2; server server-3; }}
  server {{ listen 80; location / {{ proxy_pass http://app; proxy_connect_timeout 1s; }} }}
}}"""
open("/tmp/ce-p004-nginx.conf", "w").write(conf)
sh(f"docker run -d --name lb --network {P} -p 8088:80 -v /tmp/ce-p004-nginx.conf:/etc/nginx/nginx.conf:ro nginx:1.27-alpine", check=True)
wait("curl -sf localhost:8088 >/dev/null")
loop = "for i in 1 2 3 4 5 6; do curl -s localhost:8088 | grep Hostname; done"
step("lb_rotate", loop)
step("lb_stop", "docker stop server-2")
step("lb_after", loop)
codes = "for i in $(seq 1 100); do curl -s -o /dev/null -w '%{http_code}\\n' localhost:8088; done | sort | uniq -c"
step("lb_codes", codes)
sh(f"docker start server-2")
time.sleep(2)

# ---------------------------------------------------------------- 9. load test (monitoring: p50 vs p99)
step("loadtest", "ab -n 5000 -c 50 http://lb/",
     f"docker run --rm --network {P} httpd:2.4-alpine ab -q -n 5000 -c 50 http://lb/ 2>&1 "
     "| grep -E 'Complete requests|Failed requests|Requests per second|^  50%|^  90%|^  99%|^ 100%'")

# ---------------------------------------------------------------- 3. database, index, cache
sh(f"docker run -d --name db --network {P} -e POSTGRES_PASSWORD=demo postgres:18", check=True)
sh(f"docker run -d --name redis --network {P} redis:8-alpine", check=True)
wait(f"docker exec db pg_isready -U postgres -q")
time.sleep(2)
psql = f"docker exec -i db psql -U postgres -X -q -P pager=off"
sh(f"""{psql} <<'SQL'
CREATE TABLE products (id int PRIMARY KEY, name text NOT NULL);
INSERT INTO products SELECT g, 'Product ' || g FROM generate_series(1, 1000) g;
CREATE TABLE orders (id bigserial PRIMARY KEY, customer_id int NOT NULL, product_id int NOT NULL,
                     amount numeric(10,2) NOT NULL, created_at timestamptz NOT NULL);
INSERT INTO orders (customer_id, product_id, amount, created_at)
  SELECT (random()*99999)::int + 1, (random()*999)::int + 1, round((random()*200)::numeric, 2),
         now() - (g || ' seconds')::interval
  FROM generate_series(1, 2000000) g;
VACUUM ANALYZE orders;
SQL""", check=True)
step("orders_count", "SELECT count(*) FROM orders;", f"{psql} -c 'SELECT count(*) FROM orders;'")
q = "SELECT * FROM orders WHERE customer_id = 4242;"
ea = lambda: sh(f"{psql} -c \"EXPLAIN ANALYZE {q}\" | grep -E 'Scan|Execution Time'")
ea()  # warm
step("idx_before", f"EXPLAIN ANALYZE {q}", out=ea())
step("idx_create", "CREATE INDEX orders_customer_idx ON orders (customer_id);", f"{psql} -c 'CREATE INDEX orders_customer_idx ON orders (customer_id);' && echo CREATE INDEX")
sh(f"{psql} -c 'ANALYZE orders;'")
ea()
step("idx_after", f"EXPLAIN ANALYZE {q}", out=ea())

# cache-aside in Python: same question, database vs Redis
app = r'''
import json, time, psycopg, redis
db = psycopg.connect("host=db user=postgres password=demo dbname=postgres")
r = redis.Redis(host="redis")
SQL = """SELECT p.name, sum(o.amount) AS sales FROM orders o JOIN products p ON p.id = o.product_id
         GROUP BY p.name ORDER BY sales DESC LIMIT 10"""
def top_products():
    hit = r.get("top_products")
    if hit:
        return json.loads(hit), "cache"
    rows = [[n, float(s)] for n, s in db.execute(SQL).fetchall()]
    r.set("top_products", json.dumps(rows), ex=60)   # keep it for 60 seconds (TTL)
    return rows, "database"
db.execute(SQL).fetchall()                            # warm the database's own memory first
for _ in range(3):
    t = time.perf_counter(); rows, src = top_products(); ms = (time.perf_counter() - t) * 1000
    print(f"top_products() from {src:8}  {ms:8.2f} ms")
'''.replace("%(P)s", P)
open("/tmp/ce-p004-cache.py", "w").write(app)
sh(f"docker run -d --name py --network {P} -v /tmp/ce-p004-cache.py:/cache.py:ro python:3.12-slim sleep infinity", check=True)
sh(f"docker exec py pip install -q 'psycopg[binary]==3.2.*' 'redis==5.*'", check=True)
step("cache", "python3 cache.py", f"docker exec py python /cache.py")
monitor = r'''
import statistics, time
exec(open("/cache.py").read().split("\ndb.execute(SQL)")[0])   # same top_products() as cache.py
ms = []
for i in range(1000):
    if i %% 50 == 0:
        r.delete("top_products")             # the cached copy expires: the next request goes to the database
    t = time.perf_counter(); top_products(); ms.append((time.perf_counter() - t) * 1000)
ms.sort()
print(f"requests: {len(ms)}")
print(f"average:  {statistics.mean(ms):7.2f} ms")
print(f"p50:      {ms[499]:7.2f} ms")
print(f"p99:      {ms[989]:7.2f} ms")
print(f"slowest:  {ms[-1]:7.2f} ms")
'''.replace("%%", "%")
open("/tmp/ce-p004-latency_report.py", "w").write(monitor)
sh(f"docker cp /tmp/ce-p004-latency_report.py py:/latency_report.py", check=True)
step("p99", "python3 latency_report.py", f"docker exec py python /latency_report.py")
step("cache_ttl", "redis-cli TTL top_products", f"docker exec redis redis-cli TTL top_products")

# ---------------------------------------------------------------- 4. message queue
queue = r'''
import time, redis
r = redis.Redis(host="redis")
def send_welcome_email(user):            # stand-in for a slow task: 2 s, simulated with sleep
    time.sleep(2)
def signup_direct(user):
    send_welcome_email(user)
def signup_queued(user):
    r.lpush("emails", user)               # a worker picks it up later
for name, fn in (("direct", signup_direct), ("queued", signup_queued)):
    t = time.perf_counter(); fn("ana@example.com"); ms = (time.perf_counter() - t) * 1000
    print(f"signup ({name:6}) answered in {ms:8.2f} ms")
'''.replace("%(P)s", P)
worker = r'''
import time, redis
r = redis.Redis(host="redis")
_, user = r.brpop("emails")
t = time.perf_counter(); time.sleep(2)
print(f"worker: sent welcome email to {user.decode()} ({(time.perf_counter() - t):.1f} s, in the background)")
'''.replace("%(P)s", P)
open("/tmp/ce-p004-signup.py", "w").write(queue)
open("/tmp/ce-p004-worker.py", "w").write(worker)
sh(f"docker cp /tmp/ce-p004-signup.py py:/signup.py && docker cp /tmp/ce-p004-worker.py py:/worker.py", check=True)
step("queue", "python3 signup.py", f"docker exec py python /signup.py")
step("worker", "python3 worker.py", f"docker exec py python /worker.py")

# ---------------------------------------------------------------- 5. replication + failover
pg = f"-e POSTGRES_PASSWORD=demo --network {P} postgres:18"
sh(f"docker run -d --name primary {pg}", check=True)
wait(f"docker exec primary pg_isready -U postgres -q")
time.sleep(2)
sh(f"docker exec primary bash -c 'echo \"host replication all all scram-sha-256\" >> \"$PGDATA/pg_hba.conf\"' && "
   f"docker exec primary psql -U postgres -qc 'SELECT pg_reload_conf();'", check=True)
sh(f"""docker exec -i primary psql -U postgres -X -q <<'SQL'
CREATE TABLE accounts (id int PRIMARY KEY, owner text NOT NULL, balance int NOT NULL);
INSERT INTO accounts VALUES (1, 'ana', 100), (2, 'ben', 250);
SQL""", check=True)
sh(f"docker run -d --name replica -e PGPASSWORD=demo --network {P} --entrypoint bash postgres:18 -c "
   f"'set -e; mkdir -p \"$PGDATA\"; chown postgres \"$PGDATA\"; chmod 700 \"$PGDATA\"; "
   f"gosu postgres pg_basebackup -h primary -U postgres -D \"$PGDATA\" -R -X stream; exec gosu postgres postgres'", check=True)
wait(f"docker exec replica pg_isready -U postgres -q")
time.sleep(2)
prim = f"docker exec primary psql -U postgres -X -P pager=off"
repl = f"docker exec replica psql -U postgres -X -P pager=off"
step("repl_status", "primary=# SELECT application_name, state, sync_state FROM pg_stat_replication;",
     f"{prim} -c 'SELECT application_name, state, sync_state FROM pg_stat_replication;'")
step("repl_write", "primary=# INSERT INTO accounts VALUES (3, 'cleo', 75);",
     f"{prim} -c \"INSERT INTO accounts VALUES (3, 'cleo', 75);\"")
time.sleep(0.5)
step("repl_read", "replica=# SELECT * FROM accounts;", f"{repl} -c 'SELECT * FROM accounts ORDER BY id;'")
step("repl_readonly", "replica=# INSERT INTO accounts VALUES (4, 'dev', 10);",
     f"{repl} -c \"INSERT INTO accounts VALUES (4, 'dev', 10);\"")
step("fail_stop", "docker stop primary")
step("fail_promote", "replica=# SELECT pg_promote();", f"{repl} -c 'SELECT pg_promote();'")
step("fail_write", "new primary=# INSERT INTO accounts VALUES (4, 'dev', 10);",
     f"{repl} -c \"INSERT INTO accounts VALUES (4, 'dev', 10);\" && {repl} -c 'SELECT * FROM accounts ORDER BY id;'")

# ---------------------------------------------------------------- 6. sharding
shard = r'''
import zlib
from collections import Counter
shard_of = lambda user_id: zlib.crc32(str(user_id).encode()) % 3
for uid in (17, 42, 4242, 90210):
    h = zlib.crc32(str(uid).encode())
    print(f"user {uid:>5} -> hash {h:>13,} -> % 3 = shard {h % 3}")
c = Counter(shard_of(u) for u in range(1, 1_000_001))
print("1,000,000 users:", ", ".join(f"shard {s}: {n:,}" for s, n in sorted(c.items())))
'''
open("/tmp/ce-p004-shard.py", "w").write(shard)
step("shard", "python3 shard.py", "cd /tmp && cp ce-p004-shard.py shard.py && python3 shard.py && rm shard.py")
resize = r'''
import zlib
h = lambda user_id: zlib.crc32(str(user_id).encode())
print(f"user 42: shard {h(42) % 3} of 3 -> shard {h(42) % 4} of 4")
moved = sum(h(u) % 3 != h(u) % 4 for u in range(1, 1_000_001))
print(f"going from 3 to 4 shards moves {moved:,} of 1,000,000 users ({moved / 10_000:.1f}%)")
'''
open("/tmp/ce-p004-resize.py", "w").write(resize)
step("shard_resize", "python3 resize.py", "cd /tmp && cp ce-p004-resize.py resize.py && python3 resize.py && rm resize.py")

# ---------------------------------------------------------------- 7. CDN
url = "https://cdn.jsdelivr.net/npm/react@18.3.1/umd/react.production.min.js"
sh(f"curl -s -o /dev/null {url}")
step("cdn", f"curl -s -D - -o /dev/null {url}",
     f"curl -s -D - -o /dev/null {url} | tr -d '\\r' | grep -iE '^(HTTP/|x-cache|age:|x-served-by)'")
step("cdn_time", f"curl -s -o /dev/null -w 'first byte: %{{time_starttransfer}}s' {url}",
     f"curl -s -o /dev/null -w 'first byte: %{{time_starttransfer}}s\\n' {url}")

steps["version"] = {"cmd": "versions", "out": sh(
    f"echo postgres $(docker exec replica postgres --version | awk '{{print $3}}'); "
    f"echo redis $(docker exec redis redis-server --version | awk '{{print $3}}'); "
    f"echo nginx $(docker exec lb nginx -v 2>&1 | cut -d/ -f2); date -u +%Y-%m-%dT%H:%MZ")}
json.dump(steps, open("demo_output.json", "w"), indent=1)
print("saved demo_output.json")
PY