#!/usr/bin/env python3
"""Read-only Linux Reforger memory sampler; one bounded JSONL file, no game RPCs.

Run on the server host: python3 sample_server_memory.py --output /tmp/escape-memory.jsonl
Records all Reforger processes so a host OOM can be distinguished from a container
limit. Stops after 24 hours by default. Never reads command arguments or secrets.
"""
import argparse
import datetime
import json
from pathlib import Path
import time


def fields(text):
    result = {}
    for line in text.splitlines():
        key, sep, value = line.partition(':')
        if sep and value.split() and value.split()[0].isdigit():
            result[key] = int(value.split()[0])
    return result


def snapshot(proc=Path('/proc'), cgroups=Path('/sys/fs/cgroup')):
    memory = fields((proc / 'meminfo').read_text())
    result = {'utc': datetime.datetime.now(datetime.timezone.utc).isoformat(),
              'host_kb': {k: memory.get(k) for k in ('MemTotal', 'MemAvailable', 'SwapTotal', 'SwapFree')},
              'processes': []}
    for directory in proc.iterdir():
        if not directory.name.isdigit():
            continue
        try:
            if 'ArmaReforgerServer' not in (directory / 'exe').readlink().name:
                continue
            status = fields((directory / 'status').read_text())
            # PID + kernel start tick distinguish process restarts/PID reuse.
            start = (directory / 'stat').read_text().rsplit(')', 1)[1].split()[19]
            row = {'pid': int(directory.name), 'start_ticks': int(start),
                   'kb': {k: status.get(k) for k in ('VmRSS', 'VmHWM', 'VmSize', 'RssAnon', 'RssFile', 'VmSwap')}}
            for line in (directory / 'cgroup').read_text().splitlines():
                if not line.startswith('0::'):
                    continue
                group = cgroups / line[3:].lstrip('/')
                row['cgroup'] = line[3:]
                for name in ('memory.current', 'memory.max', 'memory.peak', 'memory.events', 'memory.stat'):
                    path = group / name
                    if not path.exists():
                        continue
                    text = path.read_text().strip()
                    if name in ('memory.events', 'memory.stat'):
                        stats = dict(part.split() for part in text.splitlines())
                        if name == 'memory.stat':
                            stats = {k: v for k, v in stats.items() if k in ('anon', 'file', 'kernel', 'sock')}
                        row[name] = {k: int(v) for k, v in stats.items()}
                    else:
                        row[name] = int(text) if text.isdigit() else text
            result['processes'].append(row)
        except (OSError, ValueError, IndexError):
            # Processes may exit at any point while /proc is being read.
            continue
    return result


def append(path, record, limit=2 * 1024 * 1024):
    if path.exists() and path.stat().st_size >= limit:
        path.replace(path.with_name(path.name + '.1'))
    with path.open('a') as stream:
        stream.write(json.dumps(record, separators=(',', ':')) + '\n')


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--interval', type=int, default=60)
    parser.add_argument('--duration', type=int, default=86400)
    args = parser.parse_args()
    if args.interval < 30 or not 1 <= args.duration <= 172800:
        parser.error('interval must be >=30 seconds; duration must be 1..172800 seconds')
    deadline = time.monotonic() + args.duration
    while True:
        append(args.output, snapshot())
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            break
        time.sleep(min(args.interval, remaining))


if __name__ == '__main__':
    main()