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

import argparse
import json
import os
import sys
import time
from pathlib import Path

SENTINEL = 999999
RUNWAY_WINDOW = 180
RATE_WINDOW = 90
FLOOR_PCT = 8
METRIC_FILE = "mem_exhaustion.prom"
STATE_FILE = "mem_exhaustion.state.json"


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, ValueError, IndexError):
        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, ValueError):
        pass
    return 0.0


def runway_eta(history: list[tuple[float, int]], now: float, eff_kb: int, floor_kb: int) -> int:
    history[:] = [(ts, kb) for ts, kb in history if now - ts <= RUNWAY_WINDOW]
    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 load_history(path: Path) -> list[tuple[float, int]]:
    try:
        raw = json.loads(path.read_text())
        return [(float(ts), int(kb)) for ts, kb in raw[-64:]]
    except (OSError, ValueError, TypeError, json.JSONDecodeError):
        return []


def save_history(path: Path, history: list[tuple[float, int]]) -> None:
    try:
        path.parent.mkdir(parents=True, exist_ok=True)
        tmp = path.with_name(f"{path.name}.{os.getpid()}.tmp")
        tmp.write_text(json.dumps(history[-64:], separators=(",", ":")) + "\n")
        os.replace(tmp, path)
    except OSError:
        pass


def collect(proc_root: Path, history: list[tuple[float, int]], now: float) -> dict[str, float | int]:
    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
    history.append((now, eff))
    return {
        "node_mem_exhaustion_eta_seconds": eta,
        "node_mem_runway_percent": round((eff * 100.0 / total), 2) if total else 0.0,
        "node_mem_psi_full_avg10_percent": psi_full_avg10(proc_root),
    }


def render(metrics: dict[str, float | int]) -> str:
    lines = [
        "# HELP node_mem_exhaustion_eta_seconds Seconds until effective memory+swap reaches the configured runway floor; 999999 means no current exhaustion trajectory.",
        "# TYPE node_mem_exhaustion_eta_seconds gauge",
        f"node_mem_exhaustion_eta_seconds {metrics['node_mem_exhaustion_eta_seconds']}",
        "# HELP node_mem_runway_percent Effective memory+swap runway percentage using the existing mem_trajectory.plugin logic.",
        "# TYPE node_mem_runway_percent gauge",
        f"node_mem_runway_percent {metrics['node_mem_runway_percent']}",
        "# HELP node_mem_psi_full_avg10_percent Memory PSI full avg10 percentage from /proc/pressure/memory.",
        "# TYPE node_mem_psi_full_avg10_percent gauge",
        f"node_mem_psi_full_avg10_percent {metrics['node_mem_psi_full_avg10_percent']}",
    ]
    return "\n".join(lines) + "\n"


def write_atomic(out_dir: Path, name: str, content: str) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    tmp = out_dir / f"{name}.{os.getpid()}.tmp"
    tmp.write_text(content)
    os.replace(tmp, out_dir / name)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", default=os.environ.get("NODE_TEXTFILE_DIR", "/var/lib/node_exporter/textfile"))
    parser.add_argument("--state-dir", default=os.environ.get("SM_STATE_DIR"))
    parser.add_argument("--proc-root", default=os.environ.get("SM_PROC_ROOT", "/proc"))
    parser.add_argument("--stdout", action="store_true")
    args = parser.parse_args()

    out_dir = Path(args.output_dir)
    state_dir = Path(args.state_dir) if args.state_dir else out_dir
    history = load_history(state_dir / STATE_FILE)
    metrics = collect(Path(args.proc_root), history, time.time())
    content = render(metrics)
    if args.stdout:
        sys.stdout.write(content)
    else:
        write_atomic(out_dir, METRIC_FILE, content)
        save_history(state_dir / STATE_FILE, history)
    return 0


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