#!/usr/bin/env python3
"""bench.py — sizing bench on the N100 demo box (demo-felhom, guest 9201), through the product.

EVIDENCE, NOT PRODUCT. It presses the buttons a household presses (the endpoints the UI calls):
  GET  /api/stacks/<n>/deploy-fields · POST /api/stacks/<n>/deploy
  POST /api/stacks/<n>/stop · POST /api/stacks/<n>/remove · GET /api/stacks/<n>
Adapted from `audits/pg-last-six-2026-09-30/tools/walk.py` (same helpers, other box).

Usage:
  bench.py batch TAG app [app ...]   up + 15 min at rest + down (the light batches, unattended)
  bench.py up TAG app [app ...]      start sampling, deploy the apps one after another
  bench.py mark TAG KIND [app]       record an event (action_start / action_end ...)
  bench.py down TAG app [app ...]    15 min at rest, stop sampling, copy evidence, remove
  bench.py remove app [app ...]      remove through the product (tidy-up after an interrupted run)

While a batch runs, `guest_sampler.sh` (in the guest) and `host_sampler.sh` (on the host) record
every 5 s; their TSVs, the event log and the container->app map are copied to `evidence/<TAG>/`
BEFORE anything is removed (R-320).

The data location: this box has no registered drive, and the deploy page then says
„Nincs regisztrált adattároló — add meg kézzel az útvonalat" — so, like a household, we type one:
BENCH below, a folder under /mnt on the guest system disk (the controller sees only /mnt). It is deleted at the end of each batch.

Secrets: deploy passwords are generated per run, kept in this process (and, for the action
scripts, in a 0600 file in the scratchpad), never written to evidence.
"""
import json, os, re, secrets, subprocess, sys, time
from datetime import datetime, timezone

HERE = os.path.dirname(os.path.abspath(__file__))
AUDIT = os.path.dirname(HERE)
SC = os.environ.get("SC", "/tmp/claude-1000/-mnt-5-hdd-felhom-eu-git/73737b90-b02c-42f4-bb5f-3c922f8f6af0/scratchpad")
HOST = "felhom-pve"
GUEST = "9201"
BASE = "https://192.168.0.149"
DOMAIN = "demo-felhom.eu"
HOSTHDR = f"Host: felhom.{DOMAIN}"
BENCH = "/mnt/sizing-bench"   # the controller container sees only /mnt (guest system disk)
REST_S = int(os.environ.get("REST_S", "900"))        # 15 min at rest after the last app is up

LOG = []
EVENTS = []
GENERATED = {}       # app -> {ENV: value}; never written to evidence


def say(*a):
    line = " ".join(str(x) for x in a)
    ts = datetime.now().strftime("%H:%M:%S")
    print(f"{ts} {line}", flush=True)
    LOG.append(f"{ts} {line}")


def event(kind, app="", **kw):
    EVENTS.append(dict(ts=int(time.time()), kind=kind, app=app, **kw))


def sh(args, timeout=300, inp=None):
    try:
        return subprocess.run(args, capture_output=True, text=True, timeout=timeout, input=inp)
    except (subprocess.TimeoutExpired, OSError) as e:
        return subprocess.CompletedProcess(args, 124, "", f"{e}")


def host(script, timeout=600):
    r = sh(["ssh", "-o", "ConnectTimeout=20", HOST, "export LC_ALL=C; bash -s"], timeout=timeout, inp=script)
    return r.stdout or ""


def guest(script, timeout=600):
    """Run a bash script inside the guest. One temp file per call (guest-helper-shares-one-tmp-file)."""
    t = f"/tmp/sz{GUEST}-{os.getpid()}-{secrets.token_hex(4)}.sh"
    r = sh(["ssh", "-o", "ConnectTimeout=20", HOST,
            f"export LC_ALL=C; cat > {t}; pct push {GUEST} {t} {t} >/dev/null 2>&1; "
            f"pct exec {GUEST} -- bash {t}; pct exec {GUEST} -- rm -f {t}; rm -f {t}"],
           timeout=timeout, inp=script)
    return r.stdout or ""


def login():
    p = f"{SC}/.ctlpw"
    hdr = f"{SC}/hdr{os.getpid()}.txt"
    sh(["curl", "-sk", "-D", hdr, "-o", "/dev/null", "-H", HOSTHDR,
        "-X", "POST", "--data-urlencode", f"password@{p}", f"{BASE}/login"])
    h = open(hdr).read()
    os.remove(hdr)
    m = re.search(r"felhom_session=[A-Za-z0-9._-]+", h, re.I)
    if not m:
        sys.exit("login failed: no session cookie")
    r = sh(["curl", "-sk", "-L", "-H", HOSTHDR, "-H", f"Cookie: {m.group(0)}", f"{BASE}/"])
    c = re.search(r'<meta name="csrf-token" content="([^"]+)"', r.stdout or "")
    if not c or len(c.group(1)) != 64:
        sys.exit("login failed: no csrf token")
    return m.group(0), c.group(1)


SESS = CSRF = None


def ctl(method, path, data=None, tries=2):
    """One controller API call; re-logs in once on 302/401 (the session store is in memory)."""
    global SESS, CSRF
    if SESS is None:
        SESS, CSRF = login()
    code, body = "000", ""
    for attempt in range(tries):
        args = ["curl", "-sk", "-H", HOSTHDR, "-H", f"Cookie: {SESS}", "-w", "\n%{http_code}"]
        if method != "GET":
            args += ["-H", f"X-CSRF-Token: {CSRF}", "-H", "Content-Type: application/json", "-X", method]
            if data is not None:
                args += ["--data", json.dumps(data)]
        args.append(f"{BASE}{path}")
        r = sh(args)
        body, _, code = (r.stdout or "").rpartition("\n")
        code = code.strip()
        if code in ("302", "401") and attempt + 1 < tries:
            SESS, CSRF = login()
            continue
        break
    try:
        return code, json.loads(body)
    except Exception:
        return code, {"_raw": body[:600]}


def stack(name):
    _, d = ctl("GET", f"/api/stacks/{name}")
    return (d.get("data") or {}) if isinstance(d, dict) else {}


def deploy_values(name):
    """Fill every required deploy field as the wizard would (walk.py's rules)."""
    code, d = ctl("GET", f"/api/stacks/{name}/deploy-fields")
    fields = (((d.get("data") or {}).get("metadata") or {}).get("deploy_fields")) or []
    values = {"DOMAIN": DOMAIN}
    for f in fields:
        ev, ty = f.get("env_var"), f.get("type")
        if ev == "SUBDOMAIN":
            values[ev] = f.get("default") or name
            continue
        if ev in values:
            continue
        # `type: password` is mandatory whatever `required` says (deploy.go refuses an empty one).
        if ty == "password" or (ty == "secret" and ev in ACTION_SECRETS.get(name, ())):
            values[ev] = "Sz-" + secrets.token_hex(12)
            GENERATED.setdefault(name, {})[ev] = values[ev]
            continue
        if not f.get("required"):
            continue
        if ty == "path":
            values[ev] = BENCH
        elif ty == "secret":
            values[ev] = "Sz-" + secrets.token_hex(12)
            GENERATED.setdefault(name, {})[ev] = values[ev]
        elif f.get("default"):
            values[ev] = f["default"]
        else:
            values[ev] = f"sz-{name}"
    return values


# Optional secrets the action scripts need to know (the app's own admin login).
ACTION_SECRETS = {
    "paperless-ngx": ("PAPERLESS_ADMIN_PASSWORD",),
    "nextcloud": ("NEXTCLOUD_ADMIN_PASSWORD",),
}


def deploy(name, extra=None):
    st = stack(name)
    if st.get("deployed"):
        say(f"  {name}: already deployed — NOT touching it (someone else's app)")
        return "already"
    guest(f"mkdir -p {BENCH}")
    values = deploy_values(name)
    if extra:
        values.update(extra)
    event("deploy_start", name)
    code, d = ctl("POST", f"/api/stacks/{name}/deploy", {"values": values})
    say(f"  {name}: deploy -> {code} {str(d)[:140]}")
    if code != "202":
        event("deploy_refused", name, code=code, body=str(d)[:300])
        return "refused"
    seen = None
    for _ in range(180):                 # up to 15 min (image pulls on a home line)
        time.sleep(5)
        st = stack(name)
        seen = st.get("state")
        pins = (st.get("app_config") or {}).get("pinned_images")
        if st.get("deployed") and pins and seen in ("running", "unhealthy", "degraded"):
            event("deploy_done", name, state=seen)
            say(f"  {name}: deployed, state={seen}")
            return "ok"
        if st.get("deploy_error") or seen in ("error", "failed"):
            break
    event("deploy_failed", name, state=seen, error=str(st.get("deploy_error"))[:300])
    say(f"  {name}: never became deployed (state={seen!r}, error={str(st.get('deploy_error'))[:200]})")
    return "failed"


def remove(name):
    """Remove through the PRODUCT (live-probes rule): stop, then remove with data and backups."""
    c1, _ = ctl("POST", f"/api/stacks/{name}/stop")
    for _ in range(24):
        time.sleep(5)
        if stack(name).get("state") != "running":
            break
    code, d = ctl("POST", f"/api/stacks/{name}/remove", {"remove_hdd_data": True, "remove_backups": True})
    if code == "409":
        # R-442 fail-closed: the typed-in bench path is not a registered drive, so the product will
        # not delete data it cannot resolve. Remove the app, keep the data; the bench folder is
        # deleted by name below.
        code, d = ctl("POST", f"/api/stacks/{name}/remove", {"remove_hdd_data": False, "remove_backups": True})
    time.sleep(5)
    st = stack(name)
    left = guest(f"ls -d /opt/docker/stacks/{name}/app.yaml 2>/dev/null; "
                 f"docker ps -a --filter label=com.docker.compose.project={name} --format '{{{{.Names}}}}'; "
                 f"docker volume ls -q --filter label=com.docker.compose.project={name}")
    ok = (not st.get("deployed")) and not left.strip()
    say(f"  {name}: stop {c1}, remove {code}, deployed={st.get('deployed')}, leftovers={left.strip()!r}")
    event("removed", name, code=code, clean=ok, leftovers=left.strip()[:300])
    return ok


def samplers(tag, start):
    if start:
        host(f"rm -f /root/sizing/stop-{tag}; nohup /root/sizing/host_sampler.sh {tag} >/dev/null 2>&1 &")
        guest(f"rm -f /root/sizing/stop-{tag}; nohup /root/sizing/guest_sampler.sh {tag} >/dev/null 2>&1 &")
    else:
        host(f"touch /root/sizing/stop-{tag}")
        guest(f"touch /root/sizing/stop-{tag}")
        time.sleep(7)


def pull(tag, names):
    """Copy the evidence to DooPlex (R-320) — called BEFORE any removal."""
    d = os.path.join(AUDIT, "evidence", tag)
    os.makedirs(d, exist_ok=True)
    m = guest("docker ps -a --format '{{.Names}}\t{{.Label \"com.docker.compose.project\"}}\t{{.Image}}'")
    open(os.path.join(d, "containers.tsv"), "w").write(m)
    # the kernel's own high-water mark per container (memory.peak counts page cache too — the
    # sampled anon/used figures are the primary ones; this catches spikes shorter than 5 s)
    pk = guest("docker ps --no-trunc --format '{{.ID}} {{.Names}}' | while read id n; do "
               "c=/sys/fs/cgroup/system.slice/docker-$id.scope; "
               "echo -e \"$n\\t$(cat $c/memory.peak 2>/dev/null)\\t$(awk '$1==\"oom_kill\"{print $2}' $c/memory.events 2>/dev/null)\\t$(cat $c/memory.max 2>/dev/null)\"; done")
    open(os.path.join(d, "memory-peak.tsv"), "w").write("name\tmemory_peak_bytes\toom_kill\tmemory_max\n" + pk)
    # per-app compose limits as applied (no env values)
    lim = guest("for n in " + " ".join(names) + "; do f=/opt/docker/stacks/$n/applied-compose.yml; "
                "[ -f $f ] && echo \"== $n\" && grep -nE 'mem_limit|memory:|cpus:|image:' $f; done")
    open(os.path.join(d, "applied-limits.txt"), "w").write(lim)
    for side, cmd in (("guest", f"pct exec {GUEST} -- cat /root/sizing/guest-{tag}.tsv"),
                      ("host", f"cat /root/sizing/host-{tag}.tsv")):
        r = sh(["ssh", HOST, f"export LC_ALL=C; {cmd}"], timeout=600)
        open(os.path.join(d, f"{side}.tsv"), "w").write(r.stdout or "")
    json.dump(EVENTS, open(os.path.join(d, "events.json"), "w"), indent=1)
    open(os.path.join(d, "run.log"), "w").write("\n".join(LOG) + "\n")
    say(f"  evidence copied to {d}")


def _evfile(tag):
    d = os.path.join(AUDIT, "evidence", tag); os.makedirs(d, exist_ok=True)
    return os.path.join(d, "events.json")


def _load(tag):
    global EVENTS
    f = _evfile(tag)
    EVENTS = json.load(open(f)) if os.path.exists(f) else []


def _save(tag):
    json.dump(EVENTS, open(_evfile(tag), "w"), indent=1)


def up(tag, names):
    """Start sampling and deploy; leaves the apps running for an action."""
    _load(tag)
    samplers(tag, True)
    event("batch_start", tag=tag)
    res = {n: deploy(n) for n in names}
    event("all_up", results=res)
    _save(tag)
    if GENERATED:   # the action scripts need the app's own admin login; scratchpad only, 0600
        f = os.path.join(SC, f".generated-{tag}.json")
        fd = os.open(f, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        os.write(fd, json.dumps(GENERATED).encode()); os.close(fd)
    open(os.path.join(AUDIT, "evidence", tag, "up.log"), "a").write("\n".join(LOG) + "\n")
    return res


def mark(tag, kind, app="", **kw):
    _load(tag); event(kind, app, **kw); _save(tag)


def down(tag, names, rest=True):
    """Hold the rest window, stop sampling, copy evidence, remove through the product."""
    _load(tag)
    if rest:
        say(f"  resting {REST_S}s"); event("rest_start"); _save(tag)
        time.sleep(REST_S)
        event("rest_end"); _save(tag)
    samplers(tag, False)
    pull(tag, names)
    up_res = next((e.get("results") for e in EVENTS if e["kind"] == "all_up"), {}) or {}
    for n in names:
        if up_res.get(n) in ("ok", "failed"):
            remove(n)
    guest(f"rm -rf {BENCH}; ls -d {BENCH} 2>/dev/null || echo bench-folder-gone")
    _save(tag)
    open(os.path.join(AUDIT, "evidence", tag, "run.log"), "a").write("\n".join(LOG) + "\n")


if __name__ == "__main__":
    cmd = sys.argv[1] if len(sys.argv) > 1 else ""
    if cmd == "batch" and len(sys.argv) >= 4:
        up(sys.argv[2], sys.argv[3:]); down(sys.argv[2], sys.argv[3:])
    elif cmd == "up" and len(sys.argv) >= 4:
        print(json.dumps(up(sys.argv[2], sys.argv[3:])))
    elif cmd == "down" and len(sys.argv) >= 4:
        down(sys.argv[2], sys.argv[3:])
    elif cmd == "mark" and len(sys.argv) >= 4:
        mark(sys.argv[2], sys.argv[3], sys.argv[4] if len(sys.argv) > 4 else "")
    elif cmd == "remove" and len(sys.argv) >= 3:
        for n in sys.argv[2:]:
            remove(n)
    else:
        sys.exit(__doc__)