#!/usr/bin/env python3
"""Continuously record MAA Online memory/swap behaviour into a log.

Run once per minute by maa-online-memlog.timer. Each run appends one row to
<ROOT>/logs/memlog.csv (auto-rotated at 5 MB, one generation kept) and prints a
compact line to the journal.

Why this is instrumented the way it is: a streaming session makes the API
server's resident memory climb by tens of MB per minute and never give it back.
The interesting question each time is *where* it lands, so this sampler records

  * the service cgroup's own breakdown (anon / file / kernel / shmem),
  * the process RSS split into anonymous ([heap], [anon...], unnamed), named
    file-backed mappings, and everything else,
  * the largest named mappings, so an unexpected consumer names itself.

Anonymous memory that reaches swap is the failure mode we care about, hence the
explicit cg_swap_* and system swap columns.
"""

from __future__ import annotations

import os
import re
import subprocess
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

ROOT = Path(os.environ.get("MAA_ONLINE_ROOT") or Path(__file__).resolve().parent.parent)
LOG_DIR = Path(os.environ.get("MAA_ONLINE_LOG_DIR", ROOT / "logs"))
LOG = Path(os.environ.get("MAA_ONLINE_MEMLOG", LOG_DIR / "memlog.csv"))
MAX_BYTES = int(os.environ.get("MAA_ONLINE_MEMLOG_MAX_BYTES", 5 * 1024 * 1024))

CG = Path("/sys/fs/cgroup/system.slice/maa-online-server.service")
API = os.environ.get("MAA_ONLINE_API", "http://127.0.0.1:22888")

FIELDS = [
    "timestamp",
    "cloud_state",
    "pid",
    "proc_uptime_s",
    "cg_current_mb",
    "cg_peak_mb",
    "cg_swap_mb",
    "cg_swap_peak_mb",
    "rss_mb",
    "anon_regions",
    "big_anon_regions",
    "anon_total_mb",
    "sys_mem_avail_mb",
    "sys_swap_used_mb",
    "sys_swap_free_mb",
    "cg_anon_mb",
    "cg_file_mb",
    "cg_kernel_mb",
    "cg_shmem_mb",
    "heap_rss_mb",
    "anon_rss_mb",
    "file_rss_mb",
    "other_rss_mb",
    "top_named",
]

MIB = 1048576
ANON_NAME = re.compile(r"^\[anon")


def read_text(path: Path, default: str = "") -> str:
    try:
        return path.read_text(encoding="utf-8").strip()
    except OSError:
        return default


def cg_value(name: str) -> int:
    try:
        return int(read_text(CG / name, "0"))
    except ValueError:
        return 0


def cg_stat(name: str) -> int:
    for line in read_text(CG / "memory.stat").splitlines():
        key, _, value = line.partition(" ")
        if key == name:
            try:
                return int(value)
            except ValueError:
                return 0
    return 0


def mb(value: int) -> float:
    return round(value / MIB, 1)


def main_pid() -> int:
    try:
        out = subprocess.run(
            ["systemctl", "show", "-p", "MainPID", "--value", "maa-online-server.service"],
            capture_output=True, text=True, timeout=10, check=False,
        ).stdout.strip()
        return int(out or 0)
    except (OSError, ValueError, subprocess.SubprocessError):
        return 0


def proc_uptime(pid: int) -> int:
    stat = read_text(Path(f"/proc/{pid}/stat"))
    if not stat:
        return 0
    try:
        start_ticks = int(stat.rsplit(")", 1)[1].split()[19])
    except (IndexError, ValueError):
        return 0
    try:
        with open("/proc/uptime") as handle:
            uptime = float(handle.read().split()[0])
    except (OSError, ValueError):
        return 0
    return int(uptime - start_ticks / os.sysconf("SC_CLK_TCK"))


def proc_memory(pid: int) -> dict:
    """Split the process RSS by mapping kind, and name the biggest file mappings."""
    rss = anon_rss = file_rss = heap_rss = 0
    anon_regions = big_anon_regions = 0
    anon_total = 0
    named: dict[str, float] = {}
    name = ""
    is_anon = is_file = False
    try:
        with open(f"/proc/{pid}/smaps") as handle:
            for line in handle:
                if line[0] in "0123456789abcdef" and "-" in line.split(" ", 1)[0]:
                    parts = line.split()
                    start, end = parts[0].split("-")
                    size = (int(end, 16) - int(start, 16)) // 1024
                    name = parts[5] if len(parts) > 5 else ""
                    is_anon = name == "" or name == "[heap]" or bool(ANON_NAME.match(name))
                    is_file = name.startswith("/")
                    if is_anon:
                        anon_regions += 1
                        anon_total += size
                        if size >= 32 * 1024:
                            big_anon_regions += 1
                elif line.startswith("Rss:"):
                    value = int(line.split()[1])
                    rss += value
                    if is_anon:
                        anon_rss += value
                        if name == "[heap]":
                            heap_rss += value
                    elif is_file and name:
                        file_rss += value
                        named[name] = named.get(name, 0.0) + value / 1024
    except OSError:
        pass

    top = ", ".join(
        f"{Path(n).name}={v:.0f}MB"
        for n, v in sorted(named.items(), key=lambda kv: -kv[1])[:3]
        if v > 0
    )
    return {
        "rss_mb": round(rss / 1024, 1),
        "anon_regions": anon_regions,
        "big_anon_regions": big_anon_regions,
        "anon_total_mb": round(anon_total / 1024, 1),
        "heap_rss_mb": round(heap_rss / 1024, 1),
        "anon_rss_mb": round(anon_rss / 1024, 1),
        "file_rss_mb": round(file_rss / 1024, 1),
        "other_rss_mb": round((rss - anon_rss - file_rss) / 1024, 1),
        "top_named": top.replace(",", ";") or "-",
    }


def system_memory() -> tuple[float, float, float]:
    values: dict[str, int] = {}
    try:
        with open("/proc/meminfo") as handle:
            for line in handle:
                key, _, rest = line.partition(":")
                values[key] = int(rest.split()[0])
    except (OSError, ValueError, IndexError):
        return 0.0, 0.0, 0.0
    return (
        mb(values.get("MemAvailable", 0) * 1024),
        mb((values.get("SwapTotal", 0) - values.get("SwapFree", 0)) * 1024),
        mb(values.get("SwapFree", 0) * 1024),
    )


def cloud_state() -> str:
    try:
        with urllib.request.urlopen(f"{API}/info", timeout=5) as response:
            body = response.read().decode("utf-8", "replace")
    except (urllib.error.URLError, OSError, ValueError):
        return "unreachable"
    match = re.search(r'"status"\s*:\s*"([^"]+)"', body)
    return match.group(1) if match else "unknown"


def rotate_if_needed() -> None:
    try:
        if LOG.exists() and LOG.stat().st_size > MAX_BYTES:
            backup = LOG.with_name(LOG.name + ".1")
            backup.unlink(missing_ok=True)
            LOG.rename(backup)
    except OSError:
        pass


def main() -> int:
    LOG_DIR.mkdir(parents=True, exist_ok=True)
    rotate_if_needed()

    pid = main_pid()
    state = cloud_state()
    if pid:
        mem = proc_memory(pid)
        uptime = proc_uptime(pid)
    else:
        mem = {
            "rss_mb": 0, "anon_regions": 0, "big_anon_regions": 0, "anon_total_mb": 0,
            "heap_rss_mb": 0, "anon_rss_mb": 0, "file_rss_mb": 0, "other_rss_mb": 0,
            "top_named": "-",
        }
        uptime = 0

    mem_avail, swap_used, swap_free = system_memory()
    row = [
        time.strftime("%Y-%m-%dT%H:%M:%S%z"),
        state,
        pid,
        uptime,
        mb(cg_value("memory.current")),
        mb(cg_value("memory.peak")),
        mb(cg_value("memory.swap.current")),
        mb(cg_value("memory.swap.peak")),
        mem["rss_mb"],
        mem["anon_regions"],
        mem["big_anon_regions"],
        mem["anon_total_mb"],
        mem_avail,
        swap_used,
        swap_free,
        mb(cg_stat("anon")),
        mb(cg_stat("file")),
        mb(cg_stat("kernel")),
        mb(cg_stat("shmem")),
        mem["heap_rss_mb"],
        mem["anon_rss_mb"],
        mem["file_rss_mb"],
        mem["other_rss_mb"],
        mem["top_named"],
    ]

    fresh = not LOG.exists() or LOG.stat().st_size == 0
    with LOG.open("a", encoding="utf-8") as handle:
        if fresh:
            handle.write(",".join(FIELDS) + "\n")
        handle.write(",".join(str(value) for value in row) + "\n")

    print(
        f"{row[0]} state={state} pid={pid} cg={row[4]}MB peak={row[5]}MB cg_swap={row[6]}MB "
        f"rss={mem['rss_mb']}MB (heap={mem['heap_rss_mb']} anon={mem['anon_rss_mb']} "
        f"file={mem['file_rss_mb']} other={mem['other_rss_mb']}) "
        f"cg_anon={row[15]}MB cg_file={row[16]}MB regions={mem['anon_regions']} "
        f"big={mem['big_anon_regions']} top=[{mem['top_named']}]"
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())