/** # 脚本说明
 * 压测采样器(与发射器并行跑,同机 root):
 *   每 20s 扫描 /proc:
 *     1) 按 --user-data-dir 路径的指纹段聚合「每浏览器进程树」:
 *        进程数 / RSS 合计 / swap 合计 / CPU 累计(utime+stime, ms)
 *     2) 系统快照:MemAvailable / swap 已用 / loadavg / 浏览器树总数
 *     3) OOM 保护:非浏览器进程(且非本采样器自身)写 oom_score_adj=-600,
 *        让内核 OOM 时优先杀「易杀」的浏览器(发射器已把新浏览器主进程置 200)
 *   写 data/<label>-stats.jsonl(每轮一行 JSON)。
 *
 * 用法:bun sampler.ts --label test1 [--interval 20]
 */
const fs = require("fs") as typeof import("fs");
const args = process.argv.slice(2);
function argOf(k: string, d: string): string { const i = args.indexOf(k); return i >= 0 ? args[i + 1] : d; }
const LABEL = argOf("--label", "test1");
const INTERVAL = Number(argOf("--interval", "20")) * 1000;
const BASE = new URL(".", import.meta.url).pathname;
const OUT = `${BASE}data/${LABEL}-stats.jsonl`;

/** # 成员变量 本采样器自身进程组(不碰 oom_score_adj) **/
const SELF_PIDS = new Set([process.pid, process.ppid ?? 0]);

/** # 成员变量 每浏览器树快照类型 **/
type Tree = { fp: string; procs: number; rssKB: number; pssKB: number; swapKB: number; cpuMs: number; userDataDir: string };

/** # 函数:一轮 /proc 扫描,聚合浏览器进程树并返回系统快照 **/
function scanOnce(): { trees: Map<string, Tree>; availMB: number; swapUsedMB: number; load: number[]; cpuCount: number } {
  const trees = new Map<string, Tree>();
  const mySig = `run-stress.ts`; // 自身及其 bash 父链保护
  const txt = fs.readFileSync("/proc/meminfo", "utf8");
  const get = (k: string) => Number(new RegExp(`${k}:\\s+(\\d+)`).exec(txt)?.[1] ?? 0);
  const cpuCount = osCpus();
  const load = fs.readFileSync("/proc/loadavg", "utf8").split(/\s+/).slice(0, 3).map(Number);

  for (const pid of fs.readdirSync("/proc").filter((d) => /^\d+$/.test(d))) {
    let cmd: string;
    try { cmd = fs.readFileSync(`/proc/${pid}/cmdline`, "utf8").replace(/\0/g, " "); } catch { continue; }
    const udd = /--user-data-dir=(\S+)/.exec(cmd)?.[1] ?? null;
    if (udd) {
      // 目录形如 .../data/profiles/<fp>/browser,指纹取「browser」的父目录名
      const seg = udd.replace(/\/browser\/?$/, "").split("/").pop() ?? udd.split("/").pop() ?? "?";
      const fp = seg;
      const st = readStat(pid);
      const t = trees.get(fp) ?? { fp, procs: 0, rssKB: 0, pssKB: 0, swapKB: 0, cpuMs: 0, userDataDir: udd };
      t.procs++;
      t.rssKB += st.rssKB;
      t.pssKB += st.pssKB;
      t.swapKB += st.swapKB;
      t.cpuMs += st.cpuMs;
      trees.set(fp, t);
      // 浏览器树整体保持「易杀」(主进程已设 200,子进程继承;这里给漏网进程补救)
      try { fs.writeFileSync(`/proc/${pid}/oom_score_adj`, "200"); } catch { /* ignore */ }
    } else {
      // 非浏览器进程:OOM 保护(排除自身/内核线程写失败忽略)
      const p = Number(pid);
      if (SELF_PIDS.has(p)) continue;
      if (cmd.includes(mySig) || cmd.includes("sampler.ts") || cmd.includes("run-stress")) continue;
      try { fs.writeFileSync(`/proc/${pid}/oom_score_adj`, "-600"); } catch { /* ignore */ }
    }
  }
  const swapUsedMB = (get("SwapTotal") - get("SwapFree")) / 1024;
  const availMB = get("MemAvailable") / 1024;
  return { trees, availMB, swapUsedMB, load, cpuCount };
}

/** # 函数:读单进程 RSS/PSS/swap/CPU(kB、ms) **/
function readStat(pid: string): { rssKB: number; pssKB: number; swapKB: number; cpuMs: number } {
  let rssKB = 0, pssKB = 0, swapKB = 0, cpuMs = 0;
  try {
    const s = fs.readFileSync(`/proc/${pid}/stat`, "utf8");
    // stat 第 14/15 列为 utime/stime(tick);comm 含空格,取最后一个 ')' 之后按列号回算
    const tail = s.slice(s.lastIndexOf(")") + 2).split(" ");
    const utime = Number(tail[11]) || 0, stime = Number(tail[12]) || 0;
    cpuMs = ((utime + stime) * 1000) / 100; // CLK_TCK=100
  } catch { /* 进程已消失 */ }
  try {
    // smaps_rollup 是内核快速汇总:Pss = 独占内存(去共享重复),Swap 含交换的共享分摊
    const roll = fs.readFileSync(`/proc/${pid}/smaps_rollup`, "utf8");
    pssKB = Number(/^Pss:\s+(\d+)/m.exec(roll)?.[1] ?? 0);
    rssKB = Number(/^Rss:\s+(\d+)/m.exec(roll)?.[1] ?? 0);
    const swapPss = Number(/^SwapPss:\s+(\d+)/m.exec(roll)?.[1] ?? 0);
    const swapDirect = Number(/^Swap:\s+(\d+)/m.exec(roll)?.[1] ?? 0);
    swapKB = swapPss || swapDirect;
  } catch { /* ignore */ }
  return { rssKB, pssKB, swapKB, cpuMs };
}

/** # 函数:CPU 核数 **/
function osCpus(): number {
  try { return require("os").cpus().length; } catch { return 8; }
}

// 主循环
while (true) {
  const t0 = Date.now();
  const { trees, availMB, swapUsedMB, load, cpuCount } = scanOnce();
  const treesArr = [...trees.values()];
  const line = JSON.stringify({
    ts: new Date().toISOString(),
    browsers: treesArr.length,
    totalRssMB: +(treesArr.reduce((a, b) => a + b.rssKB, 0) / 1024).toFixed(1),
    totalPssMB: +(treesArr.reduce((a, b) => a + b.pssKB, 0) / 1024).toFixed(1),
    totalSwapMB: +(treesArr.reduce((a, b) => a + b.swapKB, 0) / 1024).toFixed(1),
    totalProcs: treesArr.reduce((a, b) => a + b.procs, 0),
    availMB: +availMB.toFixed(1), swapUsedMB: +swapUsedMB.toFixed(1),
    load, cpuCount,
    detail: treesArr,
  });
  try { fs.appendFileSync(OUT, line + "\n"); } catch { /* ignore */ }
  const wait = Math.max(0, INTERVAL - (Date.now() - t0));
  await new Promise((r) => setTimeout(r, wait));
}