// memory-lib — memory-aware admission of sub-agents (spawn_agents).
//
// Every sub-agent is a separate agent process (its own runtime and engine,
// ~150-250 MB each), and spawn_agents has no count cap: 20 or 40 children at
// once used to start together whatever the machine had left, and machines
// ran out of memory. Admission keeps the "no fixed cap" rule but starts a
// child only while the memory allows it:
//
//   - the system keeps at least MIN_FREE of its RAM available after the
//     child's expected footprint (default 20% of RAM; inside a cgroup with a
//     memory limit — a container, a systemd scope — the limit counts as RAM
//     and what it has left as available);
//   - this agent's whole process tree (the top agent, every sub-agent at
//     every nesting layer, their tools) stays under BUDGET (default 50% of
//     RAM);
//   - critically low (available under CRITICAL_FREE, default 10% of RAM):
//     nothing new starts, not even a first child, and the user is warned;
//     running children go on. Otherwise one child may always run, so a
//     batch makes progress on a small machine.
//
// The rest wait in the queue ("queued: waiting for memory") and start as
// memory frees up (a child finishing, or the next check, every second).
// Thresholds: OMNIRUSH_AGENT_MIN_FREE, OMNIRUSH_AGENT_MEMORY_BUDGET,
// OMNIRUSH_AGENT_CRITICAL_FREE and OMNIRUSH_AGENT_CHILD_ESTIMATE ("20%",
// "1.5G", "800M", a plain number is MiB), OMNIRUSH_AGENT_MEMORY=off turns
// admission off; the same keys (minFree, budget, criticalFree,
// childEstimate, enabled) under "subagentMemory" in <omnirush
// dir>/settings.json. The environment wins.
//
// Measuring: Linux reads /proc/meminfo (MemAvailable), the cgroup's
// memory.max / memory.current / memory.stat, and /proc/<pid>/statm of the
// process tree; macOS asks vm_stat and ps; elsewhere os.freemem() and an
// estimate per running child stand in for the tree.

import { execFileSync } from "node:child_process";
import { readdirSync, readFileSync } from "node:fs";
import os from "node:os";
import path from "node:path";

const MIB = 1024 * 1024;

/** Environment variable the top agent hands down: the pid whose process tree the budget covers. */
export const ENV_MEMORY_ROOT = "OMNIRUSH_AGENT_MEMORY_ROOT";

export interface MemoryPolicy {
  enabled: boolean;
  /** Keep at least this much of RAM available (bytes, resolved). */
  minFreeBytes: number;
  /** The whole process tree stays under this (bytes, resolved). */
  budgetBytes: number;
  /** Below this nothing new starts and the user is warned (bytes, resolved). */
  criticalFreeBytes: number;
  /** Footprint expected of a child before it is measured (bytes). */
  childEstimateBytes: number;
  /** Where each value came from ("env", "settings", "default"), for /agents memory. */
  sources: Record<"minFree" | "budget" | "criticalFree" | "childEstimate" | "enabled", "env" | "settings" | "default">;
}

export const DEFAULT_MIN_FREE = "20%";
export const DEFAULT_BUDGET = "50%";
export const DEFAULT_CRITICAL_FREE = "10%";
export const DEFAULT_CHILD_ESTIMATE = "256M";

/**
 * "20%" of `total`, "1.5G", "800M", "512K", or a plain number of MiB; null
 * when unreadable. Percentages are clamped to 0-100.
 */
export function parseAmount(raw: unknown, total: number): number | null {
  if (typeof raw === "number" && Number.isFinite(raw) && raw >= 0) return Math.round(raw * MIB);
  if (typeof raw !== "string") return null;
  const text = raw.trim().toLowerCase();
  if (!text) return null;
  const match = /^(\d+(?:\.\d+)?)\s*(%|b|k|kb|kib|m|mb|mib|g|gb|gib|t|tb|tib)?$/.exec(text);
  if (!match) return null;
  const value = Number(match[1]);
  const unit = match[2] ?? "m";
  if (unit === "%") return Math.round((Math.min(100, value) / 100) * total);
  const scale = unit === "b" ? 1 : unit.startsWith("k") ? 1024 : unit.startsWith("m") ? MIB : unit.startsWith("g") ? 1024 * MIB : 1024 * 1024 * MIB;
  return Math.round(value * scale);
}

/** The "subagentMemory" block of <dir>/settings.json (empty when absent). */
export function readMemorySettings(dir: string): Record<string, unknown> {
  try {
    const parsed = JSON.parse(readFileSync(path.join(dir, "settings.json"), "utf8"));
    const block = parsed?.subagentMemory;
    return block && typeof block === "object" && !Array.isArray(block) ? block : {};
  } catch {
    return {};
  }
}

/** Resolve the thresholds for a machine with `totalBytes` of RAM (env over settings over defaults). */
export function resolveMemoryPolicy(totalBytes: number, env: NodeJS.ProcessEnv = process.env, settings: Record<string, unknown> = {}): MemoryPolicy {
  const sources = {} as MemoryPolicy["sources"];
  const pick = (key: keyof MemoryPolicy["sources"], envName: string, fallback: string): number => {
    const fromEnv = parseAmount(env[envName], totalBytes);
    if (fromEnv !== null) {
      sources[key] = "env";
      return fromEnv;
    }
    const fromSettings = parseAmount(settings[key], totalBytes);
    if (fromSettings !== null) {
      sources[key] = "settings";
      return fromSettings;
    }
    sources[key] = "default";
    return parseAmount(fallback, totalBytes) ?? 0;
  };
  const rawEnabled = String(env.OMNIRUSH_AGENT_MEMORY ?? "").trim().toLowerCase();
  let enabled = true;
  if (rawEnabled) {
    sources.enabled = "env";
    enabled = !["off", "0", "false", "no", "disabled"].includes(rawEnabled);
  } else if (typeof settings.enabled === "boolean") {
    sources.enabled = "settings";
    enabled = settings.enabled;
  } else {
    sources.enabled = "default";
  }
  const minFreeBytes = pick("minFree", "OMNIRUSH_AGENT_MIN_FREE", DEFAULT_MIN_FREE);
  const budgetBytes = pick("budget", "OMNIRUSH_AGENT_MEMORY_BUDGET", DEFAULT_BUDGET);
  const criticalFreeBytes = Math.min(pick("criticalFree", "OMNIRUSH_AGENT_CRITICAL_FREE", DEFAULT_CRITICAL_FREE), minFreeBytes);
  const childEstimateBytes = Math.max(32 * MIB, pick("childEstimate", "OMNIRUSH_AGENT_CHILD_ESTIMATE", DEFAULT_CHILD_ESTIMATE));
  return { enabled, minFreeBytes, budgetBytes, criticalFreeBytes, childEstimateBytes, sources };
}

// --- measuring ---------------------------------------------------------------

export interface SystemMemory {
  /** RAM, or the cgroup's memory limit when that is smaller. */
  totalBytes: number;
  /** What can still be used before the system (or the cgroup) runs out. */
  availableBytes: number;
  source: string;
}

function readText(file: string): string | null {
  try {
    return readFileSync(file, "utf8");
  } catch {
    return null;
  }
}

/** MemTotal / MemAvailable from /proc/meminfo (bytes), or null. */
export function parseMeminfo(text: string | null): { total: number; available: number } | null {
  if (!text) return null;
  const field = (name: string) => {
    const match = new RegExp(`^${name}:\\s+(\\d+)\\s*kB`, "m").exec(text);
    return match ? Number(match[1]) * 1024 : null;
  };
  const total = field("MemTotal");
  let available = field("MemAvailable");
  if (available === null) {
    const free = field("MemFree");
    const cached = field("Cached");
    const buffers = field("Buffers");
    available = free === null ? null : free + (cached ?? 0) + (buffers ?? 0);
  }
  return total !== null && available !== null ? { total, available } : null;
}

/**
 * The tightest cgroup v2 memory limit on the way from this process's cgroup
 * to the root: {limit, used}, `used` without the page cache the kernel can
 * drop (file pages minus shmem). Null outside a limited cgroup.
 */
export function readCgroupLimit(root = "/sys/fs/cgroup", selfCgroup = readText("/proc/self/cgroup")): { limit: number; used: number; available: number } | null {
  if (!selfCgroup) return null;
  const line = selfCgroup.split("\n").find((candidate) => candidate.startsWith("0::"));
  if (!line) return null;
  let dir = line.slice(3).trim() || "/";
  let best: { limit: number; used: number; available: number } | null = null;
  for (let guard = 0; guard < 32; guard++) {
    const base = path.join(root, dir);
    const max = readText(path.join(base, "memory.max"))?.trim();
    if (max && max !== "max" && /^\d+$/.test(max)) {
      const limit = Number(max);
      const current = Number(readText(path.join(base, "memory.current"))?.trim() ?? NaN);
      const stat = readText(path.join(base, "memory.stat")) ?? "";
      const statField = (name: string) => Number(new RegExp(`^${name} (\\d+)`, "m").exec(stat)?.[1] ?? 0);
      const reclaimable = Math.max(0, statField("file") - statField("shmem"));
      const used = Number.isFinite(current) ? Math.max(0, current - reclaimable) : 0;
      const available = Math.max(0, limit - used);
      if (!best || available < best.available) best = { limit, used, available };
    }
    if (dir === "/" || dir === "") break;
    dir = path.dirname(dir);
  }
  return best;
}

/** macOS: available ≈ free + inactive + speculative + purgeable pages (vm_stat). */
export function parseVmStat(text: string): number | null {
  const pageSize = Number(/page size of (\d+) bytes/.exec(text)?.[1] ?? 4096);
  const pages = (name: string) => Number(new RegExp(`^Pages ${name}:\\s+(\\d+)`, "m").exec(text)?.[1] ?? NaN);
  const free = pages("free");
  if (!Number.isFinite(free)) return null;
  const extra = ["inactive", "speculative", "purgeable"].map(pages).filter(Number.isFinite).reduce((a, b) => a + b, 0);
  return (free + extra) * pageSize;
}

/** RAM and what is available of it right now (the cgroup's limit when tighter). */
export function readSystemMemory(): SystemMemory {
  let total = os.totalmem();
  let available = os.freemem();
  let source = "os";
  if (process.platform === "linux") {
    const info = parseMeminfo(readText("/proc/meminfo"));
    if (info) {
      total = info.total;
      available = info.available;
      source = "meminfo";
    }
    const cgroup = readCgroupLimit();
    if (cgroup && cgroup.limit < total) {
      total = cgroup.limit;
      source = "cgroup";
    }
    if (cgroup && cgroup.available < available) {
      available = cgroup.available;
      source = "cgroup";
    }
  } else if (process.platform === "darwin") {
    try {
      const parsed = parseVmStat(String(execFileSync("vm_stat", { timeout: 2_000, maxBuffer: 64 * 1024 })));
      if (parsed !== null) {
        available = Math.min(total, parsed);
        source = "vm_stat";
      }
    } catch {
      /* os.freemem() stands */
    }
  }
  return { totalBytes: total, availableBytes: Math.max(0, available), source };
}

export interface ProcessSample {
  pid: number;
  ppid: number;
  /** Private resident memory (bytes): resident minus file-backed shared pages. */
  privateBytes: number;
  /** File-backed shared resident pages (the runtime binary, mapped files). */
  sharedBytes: number;
}

/** Every process's parent and memory (Linux /proc, macOS ps); null where neither works. */
export function listProcesses(): ProcessSample[] | null {
  if (process.platform === "linux") {
    const page = 4096;
    const out: ProcessSample[] = [];
    let names: string[];
    try {
      names = readdirSync("/proc");
    } catch {
      return null;
    }
    for (const name of names) {
      if (!/^\d+$/.test(name)) continue;
      const stat = readText(`/proc/${name}/stat`);
      const statm = readText(`/proc/${name}/statm`);
      if (!stat || !statm) continue;
      // Fields after the ")" of the command name: state ppid ...
      const after = stat.slice(stat.lastIndexOf(")") + 2).split(" ");
      const ppid = Number(after[1]);
      const [, resident, shared] = statm.trim().split(/\s+/).map(Number);
      if (!Number.isFinite(ppid) || !Number.isFinite(resident)) continue;
      out.push({ pid: Number(name), ppid, privateBytes: Math.max(0, resident - (shared || 0)) * page, sharedBytes: (shared || 0) * page });
    }
    return out;
  }
  if (process.platform === "darwin") {
    try {
      const text = String(execFileSync("ps", ["-A", "-o", "pid=,ppid=,rss="], { timeout: 3_000, maxBuffer: 8 * MIB }));
      return text
        .split("\n")
        .map((line) => line.trim().split(/\s+/).map(Number))
        .filter((parts) => parts.length === 3 && parts.every(Number.isFinite))
        .map(([pid, ppid, rss]) => ({ pid, ppid, privateBytes: rss * 1024, sharedBytes: 0 }));
    } catch {
      return null;
    }
  }
  return null;
}

/**
 * Memory of `rootPid` and all its descendants: private memory summed, the
 * shared file pages counted once (the largest; every agent process maps the
 * same runtime binary). `byPid` has each member's private + shared bytes.
 */
export function treeMemory(processes: readonly ProcessSample[], rootPid: number): { bytes: number; count: number; byPid: Map<number, number> } {
  const children = new Map<number, ProcessSample[]>();
  const self = processes.find((sample) => sample.pid === rootPid);
  for (const sample of processes) {
    const list = children.get(sample.ppid);
    if (list) list.push(sample);
    else children.set(sample.ppid, [sample]);
  }
  const byPid = new Map<number, number>();
  let privateTotal = 0;
  let sharedMax = 0;
  const stack: ProcessSample[] = self ? [self] : (children.get(rootPid) ?? []).slice();
  const seen = new Set<number>();
  while (stack.length > 0) {
    const sample = stack.pop()!;
    if (seen.has(sample.pid)) continue;
    seen.add(sample.pid);
    privateTotal += sample.privateBytes;
    sharedMax = Math.max(sharedMax, sample.sharedBytes);
    byPid.set(sample.pid, sample.privateBytes + sample.sharedBytes);
    for (const child of children.get(sample.pid) ?? []) stack.push(child);
  }
  return { bytes: privateTotal + sharedMax, count: seen.size, byPid };
}

// --- the admission decision ---------------------------------------------------

export interface AdmissionInput {
  policy: MemoryPolicy;
  system: SystemMemory;
  /** The process tree's memory now (null when it cannot be measured). */
  treeBytes: number | null;
  /** Children of this process running now. */
  running: number;
  /**
   * Memory the running children have not taken yet but will (young children
   * still starting up): counted as used already.
   */
  reservedBytes: number;
  /** What the next child is expected to take. */
  estimateBytes: number;
}

export type AdmissionVerdict =
  | { ok: true }
  | { ok: false; reason: "critical" | "free" | "budget"; detail: string };

function mb(bytes: number): string {
  return bytes >= 10 * 1024 * MIB ? `${(bytes / (1024 * MIB)).toFixed(1)} GB` : `${Math.round(bytes / MIB)} MB`;
}

/** Whether one more child may start now. */
export function admit(input: AdmissionInput): AdmissionVerdict {
  const { policy, system } = input;
  if (!policy.enabled) return { ok: true };
  const available = Math.max(0, system.availableBytes - input.reservedBytes);
  if (available < policy.criticalFreeBytes) {
    return {
      ok: false,
      reason: "critical",
      detail: `memory critically low: ${mb(available)} free of ${mb(system.totalBytes)} (the floor is ${mb(policy.criticalFreeBytes)})`,
    };
  }
  const after = available - input.estimateBytes;
  if (input.running === 0) {
    // Progress: one child may run whenever the machine is not critically low.
    if (after >= policy.criticalFreeBytes) return { ok: true };
    return {
      ok: false,
      reason: "free",
      detail: `${mb(available)} free, a sub-agent needs about ${mb(input.estimateBytes)} and ${mb(policy.criticalFreeBytes)} must stay free`,
    };
  }
  if (after < policy.minFreeBytes) {
    return {
      ok: false,
      reason: "free",
      detail: `${mb(available)} free of ${mb(system.totalBytes)}; ${mb(policy.minFreeBytes)} stays free, a sub-agent needs about ${mb(input.estimateBytes)}`,
    };
  }
  if (input.treeBytes !== null && input.treeBytes + input.reservedBytes + input.estimateBytes > policy.budgetBytes) {
    return {
      ok: false,
      reason: "budget",
      detail: `agents use ${mb(input.treeBytes + input.reservedBytes)} of their ${mb(policy.budgetBytes)} budget; a sub-agent needs about ${mb(input.estimateBytes)}`,
    };
  }
  return { ok: true };
}

/** A child this process started: its pid and when, for the reservation. */
interface Tracked {
  pid: number | null;
  startedAt: number;
}

/** How long a new child counts at least its estimate (it is still loading). */
export const CHILD_WARMUP_MS = 15_000;

/**
 * The admission state of one agent process: measures (at most every
 * `sampleMs`), keeps track of its own running children, and answers
 * `check()`. Injectable readers keep it testable.
 */
export class MemoryAdmission {
  readonly policy: MemoryPolicy;
  private readonly rootPid: number;
  private readonly now: () => number;
  private readonly readSystem: () => SystemMemory;
  private readonly readProcesses: () => ProcessSample[] | null;
  private readonly sampleMs: number;
  private readonly tracked = new Map<object, Tracked>();
  private sample: { at: number; system: SystemMemory; tree: ReturnType<typeof treeMemory> | null } | null = null;
  /** Largest footprint measured of a child of this process (bytes). */
  private observedChildBytes = 0;

  constructor(options: {
    policy?: MemoryPolicy;
    env?: NodeJS.ProcessEnv;
    settings?: Record<string, unknown>;
    rootPid?: number;
    now?: () => number;
    readSystem?: () => SystemMemory;
    readProcesses?: () => ProcessSample[] | null;
    sampleMs?: number;
  } = {}) {
    const env = options.env ?? process.env;
    this.now = options.now ?? Date.now;
    this.readSystem = options.readSystem ?? readSystemMemory;
    this.readProcesses = options.readProcesses ?? listProcesses;
    this.sampleMs = options.sampleMs ?? 1_000;
    const rootFromEnv = Number(env[ENV_MEMORY_ROOT]);
    this.rootPid = options.rootPid ?? (Number.isInteger(rootFromEnv) && rootFromEnv > 0 ? rootFromEnv : process.pid);
    this.policy = options.policy ?? resolveMemoryPolicy(this.readSystem().totalBytes, env, options.settings ?? {});
  }

  /** The environment a child gets: the tree root it counts against. */
  childEnv(): Record<string, string> {
    return { [ENV_MEMORY_ROOT]: String(this.rootPid) };
  }

  private measure(force = false) {
    const at = this.now();
    if (!force && this.sample && at - this.sample.at < this.sampleMs) return this.sample;
    const system = this.readSystem();
    const processes = this.readProcesses();
    const tree = processes ? treeMemory(processes, this.rootPid) : null;
    if (tree) {
      for (const entry of this.tracked.values()) {
        const bytes = entry.pid !== null ? tree.byPid.get(entry.pid) : undefined;
        if (bytes !== undefined && at - entry.startedAt >= CHILD_WARMUP_MS) this.observedChildBytes = Math.max(this.observedChildBytes, bytes);
      }
    }
    this.sample = { at, system, tree };
    return this.sample;
  }

  /** What the next child is expected to take: the estimate, or more once children were measured bigger. */
  estimate(): number {
    return Math.max(this.policy.childEstimateBytes, Math.min(this.observedChildBytes, 4 * this.policy.childEstimateBytes));
  }

  /** Memory young children will still take (estimate minus what they use now). */
  private reserved(tree: ReturnType<typeof treeMemory> | null): number {
    const at = this.now();
    let total = 0;
    for (const entry of this.tracked.values()) {
      if (at - entry.startedAt >= CHILD_WARMUP_MS) continue;
      const used = entry.pid !== null ? tree?.byPid.get(entry.pid) ?? 0 : 0;
      total += Math.max(0, this.estimate() - used);
    }
    return total;
  }

  /** May one more child start now? */
  check(): AdmissionVerdict & { system: SystemMemory; treeBytes: number | null } {
    const sample = this.measure();
    const treeBytes = sample.tree ? sample.tree.bytes : null;
    // Without a process list (Windows), the children's estimates stand in for the tree.
    const tree = treeBytes ?? (this.tracked.size > 0 ? this.tracked.size * this.estimate() : null);
    const verdict = admit({
      policy: this.policy,
      system: sample.system,
      treeBytes: tree,
      running: this.tracked.size,
      reservedBytes: this.reserved(sample.tree),
      estimateBytes: this.estimate(),
    });
    return { ...verdict, system: sample.system, treeBytes: tree };
  }

  /** A child started (its pid once known): it counts from now on. */
  started(key: object, pid: number | null = null): void {
    // No new measurement for it: until the next sample (and while it warms
    // up) the reservation counts its estimate.
    this.tracked.set(key, { pid, startedAt: this.now() });
  }

  /** The child's pid arrived after it was started. */
  spawned(key: object, pid: number | null | undefined): void {
    const entry = this.tracked.get(key);
    if (entry && typeof pid === "number") entry.pid = pid;
  }

  /** A child ended. */
  finished(key: object): void {
    this.tracked.delete(key);
    if (this.sample) this.sample = { ...this.sample, at: -Infinity };
  }

  /** Children of this process that count now. */
  get running(): number {
    return this.tracked.size;
  }

  /** One line for /agents and status: RAM, available, the tree, the thresholds. */
  describe(): string {
    const sample = this.measure(true);
    const tree = sample.tree ? mb(sample.tree.bytes) : "unknown";
    const p = this.policy;
    if (!p.enabled) return `memory admission off (${mb(sample.system.availableBytes)} free of ${mb(sample.system.totalBytes)})`;
    return `memory: ${mb(sample.system.availableBytes)} free of ${mb(sample.system.totalBytes)} (${sample.system.source}); agents ${tree} of ${mb(p.budgetBytes)} budget; keeps ${mb(p.minFreeBytes)} free, stops at ${mb(p.criticalFreeBytes)}; ~${mb(this.estimate())} per sub-agent`;
  }
}