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

import os
import sys
import time
from collections import deque
from pathlib import Path


SENTINEL = 999999
RUNWAY_WINDOW = 180
RATE_WINDOW = 90
FLOOR_PCT = 8


def read_meminfo(proc_root: Path) -> dict[str, int]:
    values = {"MemTotal": 0, "MemAvailable": 0, "SwapTotal": 0, "SwapFree": 0}
    try:
        for line in (proc_root / "meminfo").read_text().splitlines():
            key = line.split(":", 1)[0]
            if key in values:
                values[key] = int(line.split()[1])
    except OSError:
        pass
    return values


def psi_full_avg10(proc_root: Path) -> float:
    try:
        for line in (proc_root / "pressure" / "memory").read_text().splitlines():
            if line.startswith("full "):
                for part in line.split():
                    if part.startswith("avg10="):
                        return float(part.split("=", 1)[1])
    except OSError:
        pass
    return 0.0


def runway_eta(history, now: float, eff_kb: int, floor_kb: int) -> int:
    while history and now - history[0][0] > RUNWAY_WINDOW:
        history.pop(0) if isinstance(history, list) else history.popleft()
    ref = None
    for ts, kb in history:
        if now - ts <= RATE_WINDOW:
            ref = (ts, kb)
            break
    if ref is None and history:
        ref = history[0]
    if ref is None or now - ref[0] < 20:
        return SENTINEL
    burn = (ref[1] - eff_kb) / (now - ref[0])
    if burn <= 0:
        return SENTINEL
    return int(max(0.0, eff_kb - floor_kb) / burn)


def collect(proc_root: Path = Path("/proc"), history=None, now: float | None = None) -> dict[str, float | int]:
    now = time.time() if now is None else now
    history = history if history is not None else []
    mem = read_meminfo(proc_root)
    total = mem["MemTotal"] + mem["SwapTotal"]
    eff = mem["MemAvailable"] + mem["SwapFree"]
    floor = total * FLOOR_PCT // 100 if total else 0
    eta = runway_eta(history, now, eff, floor) if total else SENTINEL
    if history is not None:
        history.append((now, eff))
    return {
        "runway_pct": round((eff * 100.0 / total), 2) if total else 0.0,
        "eta_seconds": eta,
        "psi_full_avg10": psi_full_avg10(proc_root),
    }


def emit_definition(update: int) -> None:
    print(f"CHART system_monitor.mem_runway '' 'Memory runway' 'percent seconds percent' system_monitor mem_runway line 60000 {update}")
    print("DIMENSION runway_pct 'runway_pct' absolute 1 1")
    print("DIMENSION eta_seconds 'eta_seconds' absolute 1 1")
    print("DIMENSION psi_full_avg10 'psi_full_avg10' absolute 1 1")


def emit_values(metrics: dict) -> None:
    print("BEGIN system_monitor.mem_runway")
    for key in ("runway_pct", "eta_seconds", "psi_full_avg10"):
        print(f"SET {key} = {metrics[key]}")
    print("END")
    sys.stdout.flush()


def main() -> int:
    update = 5
    proc_root = Path(os.environ.get("SM_PROC_ROOT", "/proc"))
    history = deque()
    emit_definition(update)
    while True:
        emit_values(collect(proc_root, history))
        if os.environ.get("SM_ONESHOT"):
            return 0
        time.sleep(update)


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