#!/usr/bin/env python3
from __future__ import annotations

import os
import sys
import time
from pathlib import Path


DEFAULT_BUILD = {"tsc", "vite", "esbuild", "rollup", "webpack", "turbo", "pnpm", "npm", "yarn", "cargo", "rustc", "cc", "clang", "gcc", "ld", "make", "ninja", "go", "bun"}


def total_ticks(proc_root: Path) -> int:
    try:
        parts = (proc_root / "stat").read_text().splitlines()[0].split()[1:]
        return sum(int(p) for p in parts)
    except (OSError, ValueError, IndexError):
        return 0


def parse_proc_stat(text: str) -> tuple[str, int]:
    end = text.rfind(")")
    name = text[text.find("(") + 1:end]
    parts = text[end + 2:].split()
    utime = int(parts[11]) if len(parts) > 11 else 0
    stime = int(parts[12]) if len(parts) > 12 else 0
    return name, utime + stime


def exe_name(pid_dir: Path, fallback: str) -> str:
    try:
        return os.path.basename(os.readlink(pid_dir / "exe"))
    except OSError:
        return fallback


def is_build_group(name: str, cmdline: str, cgroup: str, build_allowlist: list[str]) -> bool:
    base = name.split("/")[-1]
    if "builds.slice" in cgroup:
        return True
    if base not in set(build_allowlist):
        return False
    if base == "node":
        return any(word in cmdline for word in (" build", " dev", "vite", "webpack", "rollup", "turbo"))
    if base in {"go", "bun"}:
        return any(word in cmdline for word in (" build", " test", " run"))
    return True


def group_ticks(proc_root: Path) -> dict[str, dict[str, int | bool]]:
    groups: dict[str, dict[str, int | bool]] = {}
    for entry in proc_root.iterdir():
        if not entry.name.isdigit():
            continue
        try:
            fallback, ticks = parse_proc_stat((entry / "stat").read_text())
            name = exe_name(entry, fallback)
            cmdline = (entry / "cmdline").read_text().replace("\0", " ")
            cgroup = (entry / "cgroup").read_text()
        except OSError:
            continue
        group = groups.setdefault(name, {"ticks": 0, "build": False})
        group["ticks"] = int(group["ticks"]) + ticks
        group["build"] = bool(group["build"]) or is_build_group(name, cmdline, cgroup, list(DEFAULT_BUILD))
    return groups


def collect(proc_root: Path = Path("/proc"), previous: dict | None = None, interval: float = 10.0, build_allowlist: list[str] | None = None):
    build_allowlist = build_allowlist or list(DEFAULT_BUILD)
    groups = group_ticks(proc_root)
    current = {"total": total_ticks(proc_root), "groups": {name: int(data["ticks"]) for name, data in groups.items()}}
    if not previous:
        return {"nonbuild_busy_pct": 0.0, "build_busy_pct": 0.0}, current
    total_delta = max(1, current["total"] - int(previous.get("total", 0)))
    build = 0
    nonbuild = 0
    for name, ticks in current["groups"].items():
        delta = max(0, ticks - int(previous.get("groups", {}).get(name, 0)))
        if is_build_group(name, "", "builds.slice" if groups[name]["build"] else "", build_allowlist):
            build = max(build, delta)
        else:
            nonbuild = max(nonbuild, delta)
    scale = 100.0 / total_delta
    return {"nonbuild_busy_pct": round(nonbuild * scale, 2), "build_busy_pct": round(build * scale, 2)}, current


def main() -> int:
    update = 10
    proc_root = Path(os.environ.get("SM_PROC_ROOT", "/proc"))
    print(f"CHART system_monitor.cpu_runaway '' 'CPU runaway' 'percent' system_monitor cpu_runaway line 60020 {update}")
    print("DIMENSION nonbuild_busy_pct 'nonbuild_busy_pct' absolute 1 1")
    print("DIMENSION build_busy_pct 'build_busy_pct' absolute 1 1")
    previous = None
    while True:
        metrics, previous = collect(proc_root, previous, update)
        print("BEGIN system_monitor.cpu_runaway")
        print(f"SET nonbuild_busy_pct = {metrics['nonbuild_busy_pct']}")
        print(f"SET build_busy_pct = {metrics['build_busy_pct']}")
        print("END")
        sys.stdout.flush()
        if os.environ.get("SM_ONESHOT"):
            return 0
        time.sleep(update)


if __name__ == "__main__":
    raise SystemExit(main())