#!/usr/bin/env python3
"""Variant: N threads on one `codex app-server`, 4 KiB stdout pipe, reader stalled
while each turn produces MBs of exec output and the code-mode host is killed.

usage: driver_scale.py <outdir> <stall_seconds> <threads>
"""
import fcntl, json, os, subprocess, sys, threading, time

outdir, stall, nthreads = sys.argv[1], float(sys.argv[2]), int(sys.argv[3])
os.makedirs(outdir, exist_ok=True)
log = open(os.path.join(outdir, "events.jsonl"), "w")
t0 = time.time()
lock = threading.Lock()

def rec(kind, obj):
    with lock:
        log.write(json.dumps({"t": round(time.time() - t0, 3), "kind": kind, "msg": obj}) + "\n")
        log.flush()

proc = subprocess.Popen(["codex", "app-server"], stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                        stderr=open(os.path.join(outdir, "stderr.log"), "w"), cwd=outdir)
fcntl.fcntl(proc.stdout.fileno(), 1031, 4096)
reading = threading.Event(); reading.set()
msgs = []; cond = threading.Condition()
bytes_in = [0]

def reader():
    for line in proc.stdout:
        reading.wait()
        bytes_in[0] += len(line)
        m = json.loads(line)
        if m.get("method", "").endswith("Delta"):
            continue
        rec("in", m)
        with cond:
            msgs.append(m); cond.notify_all()
threading.Thread(target=reader, daemon=True).start()

nid = [0]
def send(method, params, notify=False):
    m = {"method": method, "params": params}
    if not notify:
        nid[0] += 1; m["id"] = nid[0]
    rec("out", m)
    proc.stdin.write((json.dumps(m) + "\n").encode()); proc.stdin.flush()
    return m.get("id")

def wait_for(pred, timeout):
    end = time.time() + timeout
    with cond:
        while True:
            for m in msgs:
                if pred(m):
                    return m
            left = end - time.time()
            if left <= 0:
                return None
            cond.wait(left)

def resp(i):
    return wait_for(lambda m: m.get("id") == i and "method" not in m, 120)

resp(send("initialize", {"clientInfo": {"name": "repro-24322", "version": "0"}}))
send("initialized", {}, notify=True)
tids = []
for _ in range(nthreads):
    r = resp(send("thread/start", {"cwd": outdir, "approvalPolicy": "never",
                                   "sandbox": "danger-full-access", "ephemeral": True}))
    tids.append(r["result"]["thread"]["id"])

def descendants(pid):
    out = []
    for task in os.listdir(f"/proc/{pid}/task"):
        try:
            kids = open(f"/proc/{pid}/task/{task}/children").read().split()
        except OSError:
            kids = []
        for k in kids:
            out.append(int(k)); out.extend(descendants(int(k)))
    return out

PROMPT = ("Use your shell/exec tool to run each of these as a separate tool call, one at a time: "
          "`seq 1 400000`, `echo one`, `seq 1 400000`, `echo two`. Then reply DONE.")

reading.clear(); rec("note", f"reader stalled for {stall}s")
for d in descendants(proc.pid):
    try:
        if os.readlink(f"/proc/{d}/exe").endswith("/codex-code-mode-host"):
            os.kill(d, 9); rec("note", f"killed code-mode host pid {d}")
    except OSError:
        pass
for tid in tids:
    send("turn/start", {"threadId": tid, "input": [{"type": "text", "text": PROMPT}]})
time.sleep(stall)
rec("note", "reader resumed"); reading.set()
for tid in tids:
    done = wait_for(lambda m: m.get("method") == "turn/completed" and m["params"].get("threadId") == tid, 900)
    rec("note", f"thread {tid} completed={done is not None}")
rec("note", f"bytes read from app-server stdout: {bytes_in[0]}")
proc.terminate(); proc.wait(10)
rec("note", "exit")