#!/usr/bin/env python3
"""wait_shells.py <main-executable-path> <count> <timeout-s>

Wait until <count> shells (zsh/bash, including "-zsh" login shells and shells
started through /usr/bin/login) run below the app's main process (the process
whose executable is <main-executable-path>), then print how long it took."""
import subprocess, sys, time

exe, want, timeout = sys.argv[1], int(sys.argv[2]), float(sys.argv[3])
t0, shells = time.time(), []
while time.time() - t0 < timeout:
    rows = subprocess.run(["ps", "-axo", "pid=,ppid=,comm="], capture_output=True, text=True).stdout.splitlines()
    procs = {}
    for r in rows:
        f = r.split(None, 2)
        if len(f) == 3:
            procs[int(f[0])] = (int(f[1]), f[2])
    mains = {p for p, (pp, c) in procs.items() if c == exe}

    def under_main(p):
        seen = set()
        while p in procs and p not in seen:
            seen.add(p)
            p = procs[p][0]
            if p in mains:
                return True
        return False

    shells = [p for p, (pp, c) in procs.items()
              if c.split("/")[-1].lstrip("-") in ("zsh", "bash") and under_main(p)]
    if mains and len(shells) >= want:
        print(f"{len(shells)} shells after {time.time() - t0:.1f} s ({time.strftime('%H:%M:%S')})")
        sys.exit(0)
    time.sleep(0.25)
print(f"timeout: {len(shells)} shells")
sys.exit(1)