#!/usr/bin/env python3
# Persistent-connection AUTH stress test for CUBELinux-2 cube-server.
#
# Server's REAL model: one signed-HELLO R4 handshake per connection
# (auth-once), then unlimited ops over the open socket (no per-command
# re-auth). Fair comparison vs etcd/Redis (auth once at connect).
#
# Wire: 4-byte little-endian length prefix + UTF-8 payload (cubesys/src/net.rs).
# Handshake: server "CHALLENGE <nonce>"; client "HELLO <tenant> <owner> <sig>"
#   sig = HMAC-SHA256(psk, "nonce|tenant|owner_local|") hex (owner_remote absent).
#
# Socket has a 3s recv timeout so a WAL-stalled daemon cannot hang the test
# forever; timed-out ops are counted as failures (revealing the sustainable
# rate under load). ALL artifacts live on real disk (/root/cube-stress), never
# tmpfs, because the WAL can exceed tmpfs capacity.
import os, sys, socket, struct, time, threading, csv, hmac, hashlib, subprocess, tempfile, shutil
from collections import defaultdict
import statistics as st

REPO = os.environ.get("REPO", "/home/CUBELinux/CUBELinux-2")
SRV = os.path.join(REPO, "target/release/cube-server")
if not os.path.exists(SRV):
    sys.exit("cube-server release binary missing; build first")

USERS = int(os.environ.get("USERS", "8"))
DURATION = int(os.environ.get("DURATION", "150"))
PER_TENANT = os.environ.get("PER_TENANT", "0") == "1"
PSK = b"lan-shared-key-1234"

WORK = "/root/cube-stress"
os.makedirs(WORK, exist_ok=True)
TMPD = tempfile.mkdtemp(prefix=f"{WORK}/run-")
STORE = os.path.join(TMPD, f"store-{os.getpid()}.json")
PSK_FILE = os.path.join(TMPD, "psk.txt")
open(PSK_FILE, "w").write(PSK.decode())
SOCK = os.path.join(TMPD, "cube.sock")
RUN_LOG = os.path.join(TMPD, "run.log")
DAEMON_LOG = os.path.join(TMPD, "daemon.log")
STATS_LOG = os.path.join(TMPD, "stats-samples.log")
ERR_LOG = os.path.join(TMPD, "run-errors.log")
CSV = os.path.join(TMPD, "cmds.csv")

TENANTS = ["alpha","bravo","charlie","delta","echo","foxtrot","golf","hotel","india","juliet"]
OWNERS  = ["alice","bob","carol","dave","erin","frank","grace","heidi","ivan","judy"]

def log(*a):
    with open(RUN_LOG, "a") as f:
        f.write(" ".join(map(str,a)) + "\n")
log(f"USERS={USERS} DURATION={DURATION}s PER_TENANT={PER_TENANT} MODE=persistent-connection")

def write_frame(sock, payload):
    b = payload.encode(); sock.sendall(struct.pack("<I", len(b)) + b)
def read_frame(sock):
    n = struct.unpack("<I", sock.recv(4))[0]
    return sock.recv(n).decode()
def sign_hello(psk, nonce, tenant, owner_local, owner_remote=""):
    msg = f"{nonce}|{tenant}|{owner_local}|{owner_remote}".encode()
    return hmac.new(psk, msg, hashlib.sha256).hexdigest()

daemon_args = [SRV, "--socket", SOCK, "--store", STORE, "--auth-key", PSK_FILE]
if PER_TENANT:
    tdir = os.path.join(TMPD, "tenants"); os.makedirs(tdir)
    daemon_args += ["--tenant-dir", tdir]
dp = subprocess.Popen(daemon_args, stdout=open(DAEMON_LOG,"w"), stderr=subprocess.STDOUT)
for _ in range(15):
    if os.path.exists(SOCK): break
    time.sleep(1)
if not os.path.exists(SOCK):
    log("daemon failed to start"); sys.exit(1)
log(f"daemon up pid={dp.pid} on {SOCK} [auth REQUIRED, persistent connections]")

results = []
stop = threading.Event()
csv_lock = threading.Lock()

def make_cmd(u, c):
    pair = c // 8
    X = ((pair + u*7) % 200) + 1
    Y = ((pair*3 + u*5) % 200) + 1
    Z = ((pair*11 + u*13) % 200) + 1
    W = ((pair*17 + u*3) % 200) + 1
    p = f"/c{W:03d}/z{X:03d}/y{Y:03d}/x{Z:03d}"
    op = c % 8
    if op == 0:  return op, f"prog {p} const {X} halt"
    if op == 1:  return op, f"run {p}"
    if op == 2:  return op, f"grant {OWNERS[u]} r 100.1.1.1"
    if op == 3:  return op, f"revoke {OWNERS[u]} 100.1.1.1"
    if op == 4:  return op, "query fn"
    if op == 5:  return op, "stats"
    if op == 6:
        if os.environ.get("NO_AUDIT") == "1":
            return op, "stats"   # drop the slow audit op to replicate legacy op mix
        return op, "audit"
    return op, "stats"

def drive_user(u):
    tenant, owner = TENANTS[u], OWNERS[u]
    s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM); s.connect(SOCK)
    s.settimeout(3.0)
    errs_logged = 0
    t0 = time.time()
    chal = read_frame(s); assert chal.lower().startswith("challenge ")
    nonce = chal[len("challenge "):].strip()
    sig = sign_hello(PSK, nonce, tenant, owner)
    write_frame(s, f"HELLO {tenant} {owner} {sig}")
    ack = read_frame(s)
    hs_ms = (time.time() - t0) * 1000.0
    with csv_lock:
        results.append((u, -1, 1 if ack.startswith("ok:") else 0, 0, hs_ms))
    log(f"user {u} handshake ok={ack.startswith('ok:')} {hs_ms:.2f}ms ack={ack!r}")
    c = 0
    while not stop.is_set() and time.time() < END:
        op, cmd = make_cmd(u, c)
        t0 = time.time()
        try:
            write_frame(s, cmd); reply = read_frame(s)
            ok = 0 if reply.lower().startswith("error") else 1
            if ok == 0 and op == 1 and errs_logged < 8:
                with open(ERR_LOG, "a") as ef:
                    ef.write(f"RUN FAIL u={u} c={c} cmd={cmd!r} reply={reply[:160]!r}\n")
                errs_logged += 1
        except socket.timeout:
            ok = 0; reply = "TIMEOUT(>3s)"
        except Exception as e:
            ok = 0; reply = f"EXC:{e}"
        lat = (time.time() - t0) * 1e9
        with csv_lock:
            results.append((u, op, ok, lat, 0.0))
        c += 1
    s.close()

END = time.time() + DURATION
threads = [threading.Thread(target=drive_user, args=(u,), daemon=True) for u in range(USERS)]
for t in threads: t.start()

def sampler():
    s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM); s.connect(SOCK)
    s.settimeout(3.0)
    chal = read_frame(s); nonce = chal[len("challenge "):].strip()
    sig = sign_hello(PSK, nonce, "alpha", "sampler")
    write_frame(s, f"HELLO alpha sampler {sig}"); read_frame(s)
    n = 0
    while not stop.is_set():
        ts = time.strftime("%H:%M:%S", time.gmtime())
        with open(STATS_LOG, "a") as f:
            f.write(f"===== sample {n} @ {ts} =====\n")
            try:
                write_frame(s, "stats"); f.write(read_frame(s) + "\n")
            except Exception as e:
                f.write(f"(sampler failed: {e})\n")
        n += 1; time.sleep(10)
    s.close()
sthread = threading.Thread(target=sampler, daemon=True); sthread.start()

for t in threads: t.join()
stop.set(); sthread.join(timeout=3)
dp.terminate()
try: dp.wait(timeout=5)
except Exception: dp.kill()

with open(CSV, "w", newline="") as f:
    w = csv.writer(f); w.writerow(["user","op","ok","latency_ns","handshake_ms"])
    for r in results: w.writerow(r)

rows = results
hs = [r[4] for r in rows if r[1] == -1]
ops = [r for r in rows if r[1] != -1]
ok = sum(1 for r in ops if r[2] == 1)
lat = [r[3] for r in ops]
def pct(xs,p):
    xs=sorted(xs); k=int(round((p/100)*(len(xs)-1))); return xs[k]/1e6
log("===== AGGREGATE RESULTS (persistent-connection, auth-once) =====")
log(f"users                : {USERS}")
log(f"duration (s)         : {DURATION}")
log(f"handshakes           : {len(hs)}  (1 per connection)")
if hs: log(f"handshake ms mean/max: {st.mean(hs):.3f} / {max(hs):.3f}")
log(f"total ops            : {len(ops)}")
log(f"successful (ok=1)    : {ok}  ({100*ok/max(len(ops),1):.2f}%)")
log(f"failed (ok=0)        : {len(ops)-ok}")
if lat:
    log(f"op latency ms mean/med : {st.mean(lat)/1e6:.3f} / {st.median(lat)/1e6:.3f}")
    log(f"op latency ms min/p50  : {min(lat)/1e6:.3f} / {pct(lat,50):.3f}")
    log(f"op latency ms p95/p99  : {pct(lat,95):.3f} / {pct(lat,99):.3f}")
    log(f"op latency ms max      : {max(lat)/1e6:.3f}")
byop=defaultdict(list); byop_ok=defaultdict(int)
for r in ops: byop[r[1]].append(r[3]); byop_ok[r[1]]+=r[2]
log("per-op (op: count, ok, mean_ms, max_ms):")
for op in sorted(byop):
    xs=byop[op]; log(f"  op{op}: n={len(xs)} ok={byop_ok[op]} mean={st.mean(xs)/1e6:.3f} max={max(xs)/1e6:.3f}")
byu=defaultdict(list); byu_ok=defaultdict(int)
for r in ops: byu[r[0]].append(r[3]); byu_ok[r[0]]+=r[2]
log("per-user (user: ops, ok, mean_ms):")
for u in sorted(byu, key=lambda x:int(x)):
    xs=byu[u]; log(f"  u{u}: ops={len(xs)} ok={byu_ok[u]} mean={st.mean(xs)/1e6:.3f}")
# Clean store/WAL so real disk doesn't accumulate; keep logs/csv.
for f in (STORE, STORE+".wal", STORE+".delta", STORE+".seq",
          os.path.expanduser("~/.local/state/cube/cube-store.recovery.ndjson")):
    try: os.remove(f)
    except OSError: pass
shutil.rmtree(os.path.join(TMPD,"tenants"), ignore_errors=True)
log(f"TMPD={TMPD}")
log("DONE.")
print(f"DONE. results in {TMPD}")