#!/usr/bin/env python3
"""Idle memory / CPU / wakeup sampler for one desktop app on macOS.

Which processes count as "the app":
  * every process whose executable lives inside the app bundle (--bundle), and
  * every process descended from one of those (helpers, renderers, extension
    hosts, the JVM's helpers, language servers, git started by the app ...),
  EXCEPT
  * terminal shells (zsh, bash, sh, fish, login, ...) and everything below them:
    whatever runs inside a terminal tab (an agent, a build) is the same in any
    host app and is reported separately, never counted;
  * processes started from the bundle with "--cli" (Next Term's MCP helper,
    which agents' MCP clients start; it is a child of the agent, not the app).
  Processes that macOS attributes to the app (its "responsible" pid), such as
  a system XPC service started for it or a daemon it launched (JetBrains'
  jetbrainsd), ARE counted, as the app caused them, unless they are a shell or
  run below one. Anything else attributed to the app is listed separately.

Numbers come from scripts/pstat (libproc, read-only, no root):
  memory = sum of phys_footprint ("Memory" in Activity Monitor) of the app's
           processes; median of --samples samples taken --interval s apart.
  CPU    = (CPU time used by the app's processes, plus CPU of short-lived
           children they reaped, during --cpu-window s) / wall time.
           100% = one core fully busy.
  wakeups/s = package-idle + interrupt wakeups over the same window.
After the CPU window, /usr/bin/footprint is run once on the same pids as a
cross-check (it is not run during the window so it cannot disturb it).

Usage:
  measure.py --label cursor --bundle "/Applications/Cursor.app/" --out raw/cursor.json
"""
import argparse, datetime, json, os, re, statistics, subprocess, sys, time

HERE = os.path.dirname(os.path.abspath(__file__))
PSTAT = os.environ.get("PSTAT", os.path.join(HERE, "pstat"))
SHELLS = {"zsh", "bash", "sh", "fish", "login", "tcsh", "csh", "nu", "dash", "ksh"}
HOME = os.path.expanduser("~")
SCRUB = []  # (needle, replacement) pairs applied to every string we write
TMPDIR = os.environ.get("TMPDIR", "").rstrip("/")


def scrub(s):
    if not isinstance(s, str):
        return s
    for a, b in SCRUB:
        s = s.replace(a, b)
    return s


def pstat_all(argmap):
    out = subprocess.run([PSTAT], capture_output=True, text=True).stdout
    procs = {}
    for line in out.splitlines():
        try:
            p = json.loads(line)
        except ValueError:
            continue
        if not p.get("path"):
            # proc_pidpath is empty when the binary was replaced on disk after the
            # process started (an app update); fall back to argv[0] then.
            args = argmap.get(p["pid"], (0, ""))[1]
            m = re.match(r"(/.*?\.app/Contents/\S*|\S+)", args)
            p["path"] = m.group(1) if m else ""
        procs[p["pid"]] = p
    for pid, (ppid, uid) in PS_TREE.items():
        if pid not in procs:  # visible to ps but not to libproc: keep it in the tree, no metrics
            args = argmap.get(pid, (0, ""))[1]
            exe = args.split(" ")[0] if args else ""
            procs[pid] = {"pid": pid, "ppid": ppid, "uid": uid, "rpid": None, "path": exe,
                          "comm": os.path.basename(exe), "unreadable": True}
    return procs


PS_TREE = {}  # pid -> (ppid, uid), from ps, for processes libproc cannot read (e.g. root-owned /usr/bin/login)


def ps_args():
    out = subprocess.run(["ps", "-axww", "-o", "pid=,ppid=,uid=,rss=,args="], capture_output=True, text=True).stdout
    res = {}
    PS_TREE.clear()
    for line in out.splitlines():
        f = line.split(None, 4)
        if len(f) >= 4:
            res[int(f[0])] = (int(f[3]), f[4] if len(f) > 4 else "")
            PS_TREE[int(f[0])] = (int(f[1]), int(f[2]))
    return res


def base(p):
    n = os.path.basename(p.get("path") or "") or p.get("comm", "")
    return n.lstrip("-")


def kind_of(p, args):
    if ".xpc/Contents/MacOS/" in (p.get("path") or "") and not (p.get("path") or "").startswith("/Applications/"):
        return "system-xpc:" + base(p)
    if p.get("ppid") == 1 and p.get("rpid") and p.get("rpid") != p["pid"] and not (p.get("path") or "").startswith("/Applications/"):
        return "attributed:" + base(p)
    m = re.search(r"--type=(\S+)", args)
    if m:
        t = m.group(1)
        m2 = re.search(r"--utility-sub-type=(\S+)", args)
        if m2:
            t += ":" + m2.group(1).split(".")[-1]
        return t
    for key in ("ptyHost", "fileWatcher", "extensionHost", "sharedProcess"):
        if key in args:
            return key
    return base(p)


def classify(procs, argmap, bundle):
    children = {}
    for p in procs.values():
        children.setdefault(p["ppid"], []).append(p["pid"])

    def subtree(pid):
        out, st = [], [pid]
        while st:
            q = st.pop()
            out.append(q)
            st.extend(children.get(q, []))
        return out

    seeds, separate = set(), set()
    for p in procs.values():
        if p["uid"] != os.getuid() or not (p.get("path") or "").startswith(bundle):
            continue
        args = argmap.get(p["pid"], (0, ""))[1]
        if re.search(r"(^|\s)--cli(\s|$)", args):
            separate.add(p["pid"])
        else:
            seeds.add(p["pid"])
    # seeds that descend from another seed are fine; separate helpers' subtrees stay separate
    app, shells = set(), set()
    for s in seeds:
        st = [s]
        while st:
            q = st.pop()
            if q in app or q in shells or q in separate:
                continue
            p = procs.get(q)
            if p is None:
                continue
            if q not in seeds and base(p) in SHELLS:
                shells.update(subtree(q))
                continue
            app.add(q)
            st.extend(children.get(q, []))
    attributed = {p["pid"] for p in procs.values()
                  if p.get("rpid") in seeds and p["pid"] not in app | shells | separate}
    # Count what macOS attributes to the app (XPC services, daemons it launched such as
    # JetBrains' jetbrainsd), unless it is a shell or below one (whatever a tab started).
    def below_shell(q):
        seen = set()
        while q and q not in seen:
            seen.add(q)
            p = procs.get(q)
            if p is None:
                return False
            if base(p) in SHELLS:
                return True
            q = p["ppid"]
        return False
    counted = {q for q in attributed if not below_shell(q)}
    return app | counted, shells, separate, attributed - counted


def row(p, argmap):
    rss_kb, args = argmap.get(p["pid"], (0, ""))
    return {"pid": p["pid"], "ppid": p["ppid"], "rpid": p.get("rpid"), "kind": scrub(kind_of(p, args)),
            "path": scrub(p.get("path")), "args": scrub(args[:400]),
            "footprint_b": p.get("footprint"), "resident_b": p.get("resident"),
            "lifetime_max_footprint_b": p.get("lifetime_max_footprint"),
            "cpu_ns": p.get("cpu_ns"), "child_cpu_ns": p.get("child_cpu_ns"),
            "wakeups": p.get("wakeups"), "child_wakeups": p.get("child_wakeups"),
            "ps_rss_kb": rss_kb}


def pressure():
    try:
        lvl = int(subprocess.run(["sysctl", "-n", "kern.memorystatus_vm_pressure_level"],
                                 capture_output=True, text=True).stdout.strip())
        return {1: "normal", 2: "warning", 4: "critical"}.get(lvl, str(lvl))
    except ValueError:
        return "unknown"


def snapshot(bundle):
    t = time.time()
    argmap = ps_args()
    procs = pstat_all(argmap)
    app, shells, separate, attributed = classify(procs, argmap, bundle)
    rows = lambda s: [row(procs[q], argmap) for q in sorted(s) if q in procs]
    a = rows(app)
    return {
        "time": datetime.datetime.fromtimestamp(t).isoformat(timespec="seconds"),
        "epoch": t, "loadavg": [round(x, 2) for x in os.getloadavg()], "memory_pressure": pressure(),
        "app_process_count": len(a),
        "app_footprint_b": sum(r["footprint_b"] or 0 for r in a),
        "app_resident_b": sum(r["resident_b"] or 0 for r in a),
        "app_ps_rss_kb": sum(r["ps_rss_kb"] for r in a),
        "app_processes": a,
        "excluded_shell_subtrees": [{k: r[k] for k in ("pid", "ppid", "kind", "path", "footprint_b")} for r in rows(shells)],
        "separate_cli_helpers": rows(separate),
        "attributed_by_macos_not_counted": rows(attributed),
    }


def footprint_tool(pids):
    if not pids:
        return {}
    cmd = ["footprint", "-f", "bytes", "--noCategories"]
    for pid in sorted(pids):
        cmd += ["-p", str(pid)]
    out = subprocess.run(cmd, capture_output=True, text=True).stdout
    res = {}
    for line in out.splitlines():
        m = re.search(r"\[(\d+)\]:.*Footprint:\s+(\d+) B", line)
        if m:
            res[int(m.group(1))] = int(m.group(2))
    return res


def wait_for_load(limit, max_wait):
    if limit <= 0:
        return None
    t0 = time.time()
    while os.getloadavg()[0] > limit and time.time() - t0 < max_wait:
        time.sleep(5)
    return round(time.time() - t0, 1)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--label", required=True)
    ap.add_argument("--bundle", required=True, help="app bundle path prefix, with trailing slash")
    ap.add_argument("--settle", type=float, default=60)
    ap.add_argument("--samples", type=int, default=3)
    ap.add_argument("--interval", type=float, default=20)
    ap.add_argument("--cpu-window", type=float, default=60)
    ap.add_argument("--wait-load", type=float, default=0, help="before settling, wait until 1-min load average <= this")
    ap.add_argument("--wait-load-max", type=float, default=600)
    ap.add_argument("--scrub", action="append", default=[], help="NEEDLE=REPLACEMENT for paths in the output")
    ap.add_argument("--note", default="")
    ap.add_argument("--out", required=True)
    a = ap.parse_args()
    for s in a.scrub:
        k, v = s.split("=", 1)
        SCRUB.append((k, v))
    if TMPDIR:
        SCRUB.append(("/private" + TMPDIR, "$TMPDIR"))
        SCRUB.append((TMPDIR, "$TMPDIR"))
    SCRUB.append((HOME, "~"))

    started = datetime.datetime.now().isoformat(timespec="seconds")
    waited = wait_for_load(a.wait_load, a.wait_load_max)
    print(f"[{a.label}] load wait {waited}s, settling {a.settle:.0f}s ...", flush=True)
    time.sleep(a.settle)
    samples, t0 = [], None
    for i in range(a.samples):
        if i:
            time.sleep(max(0, t0 + i * a.interval - time.time()))
        s = snapshot(a.bundle)
        if t0 is None:
            t0 = s["epoch"]
        samples.append(s)
        print(f"[{a.label}] sample {i+1}: {s['app_process_count']} procs, footprint "
              f"{s['app_footprint_b']/1048576:.1f} MiB, rss {s['app_resident_b']/1048576:.1f} MiB, "
              f"load {s['loadavg'][0]}, shells excluded {len(s['excluded_shell_subtrees'])}", flush=True)
    time.sleep(max(0, t0 + a.cpu_window - time.time()))
    end = snapshot(a.bundle)
    wall = end["epoch"] - t0

    first = {r["pid"]: r for r in samples[0]["app_processes"]}
    last = {r["pid"]: r for r in end["app_processes"]}
    per, cpu_ns, wk = [], 0, 0
    for pid, r in last.items():
        b = first.get(pid)
        d_cpu = (r["cpu_ns"] or 0) + (r["child_cpu_ns"] or 0)
        d_wk = (r["wakeups"] or 0) + (r["child_wakeups"] or 0)
        if b:  # existed for the whole window: delta; new during the window: all of it
            d_cpu -= (b["cpu_ns"] or 0) + (b["child_cpu_ns"] or 0)
            d_wk -= (b["wakeups"] or 0) + (b["child_wakeups"] or 0)
        d_cpu, d_wk = max(0, d_cpu), max(0, d_wk)
        cpu_ns += d_cpu
        wk += d_wk
        per.append({"pid": pid, "kind": r["kind"], "cpu_ms": round(d_cpu / 1e6, 1), "wakeups": d_wk,
                    "footprint_mib": round((r["footprint_b"] or 0) / 1048576, 1)})
    gone = sorted(set(first) - set(last))

    fp_tool = footprint_tool(set(last))
    fps = [s["app_footprint_b"] for s in samples]
    res = [s["app_resident_b"] for s in samples]
    result = {
        "label": a.label, "bundle": scrub(a.bundle), "note": a.note, "started": started,
        "load_wait_s": waited, "settle_s": a.settle, "interval_s": a.interval,
        "cpu_window_s": round(wall, 2),
        "median_app_footprint_b": statistics.median(fps),
        "median_app_footprint_mib": round(statistics.median(fps) / 1048576, 1),
        "app_footprint_mib_per_sample": [round(x / 1048576, 1) for x in fps],
        "median_app_resident_mib": round(statistics.median(res) / 1048576, 1),
        "app_process_count_per_sample": [s["app_process_count"] for s in samples] + [end["app_process_count"]],
        "avg_cpu_percent": round(100 * cpu_ns / 1e9 / wall, 3),
        "wakeups_per_s": round(wk / wall, 1),
        "processes_gone_during_window": gone,
        "footprint_tool_crosscheck_mib": round(sum(fp_tool.values()) / 1048576, 1),
        "footprint_tool_per_pid_b": fp_tool,
        "per_process_window": sorted(per, key=lambda x: -x["footprint_mib"]),
        "samples": samples,
        "end_snapshot": end,
    }
    os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True)
    with open(a.out, "w") as f:
        json.dump(result, f, indent=1)
    print(f"[{a.label}] MEDIAN footprint {result['median_app_footprint_mib']} MiB "
          f"(footprint tool {result['footprint_tool_crosscheck_mib']} MiB, resident {result['median_app_resident_mib']} MiB), "
          f"procs {result['app_process_count_per_sample']}, avg CPU {result['avg_cpu_percent']}%, "
          f"{result['wakeups_per_s']} wakeups/s over {wall:.0f}s -> {a.out}", flush=True)


if __name__ == "__main__":
    main()