#!/usr/bin/env python3
"""Stability soak — sample backend health, RSS, and open sockets over time.

Runs periodic probes against a live mugi_backend and fails if:
  - /health stops responding
  - backend RSS grows beyond a bounded slope (leak signal)
  - open TCP sockets to the backend port grow without bound

Usage:
  ~/.mugi/venv/bin/python scripts/soak.py --hours 2
  ~/.mugi/venv/bin/python scripts/soak.py --minutes 30 --interval 60 --json /tmp/soak.json

Requires a listening backend (``make backend`` or Swift app).
"""
from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import time
import urllib.error
import urllib.request
from typing import Any

_DEFAULT_PORT = 8321
_DEFAULT_INTERVAL = 120.0
_RSS_SLOPE_MB_PER_HOUR = 200.0  # fail if RSS grows faster than this
_SOCKET_GROWTH_MAX = 8  # fail if socket count grows by more than this vs baseline


def _health_ok(port: int) -> bool:
    try:
        with urllib.request.urlopen(f"http://127.0.0.1:{port}/health", timeout=5) as resp:
            return resp.status == 200
    except (urllib.error.URLError, TimeoutError, OSError):
        return False


def _backend_pid(port: int) -> int | None:
    try:
        out = subprocess.check_output(
            ["lsof", "-ti", f":{port}"],
            text=True,
            stderr=subprocess.DEVNULL,
        ).strip()
        if not out:
            return None
        return int(out.split("\n")[0])
    except (subprocess.CalledProcessError, ValueError, FileNotFoundError):
        return None


def _rss_mb(pid: int) -> float | None:
    try:
        out = subprocess.check_output(["ps", "-o", "rss=", "-p", str(pid)], text=True)
        kb = int(out.strip())
        return kb / 1024.0
    except (subprocess.CalledProcessError, ValueError, FileNotFoundError):
        return None


def _socket_count(port: int) -> int:
    try:
        out = subprocess.check_output(
            ["lsof", "-i", f":{port}"],
            text=True,
            stderr=subprocess.DEVNULL,
        )
        return max(0, len([ln for ln in out.splitlines() if ln.strip()]) - 1)
    except (subprocess.CalledProcessError, FileNotFoundError):
        return 0


def main() -> int:
    parser = argparse.ArgumentParser(description="Mugi backend stability soak")
    parser.add_argument("--hours", type=float, default=2.0, help="Soak duration (default 2h)")
    parser.add_argument("--minutes", type=float, default=0.0, help="Override duration in minutes")
    parser.add_argument("--interval", type=float, default=_DEFAULT_INTERVAL, help="Sample interval seconds")
    parser.add_argument("--port", type=int, default=_DEFAULT_PORT)
    parser.add_argument("--json", default="", help="Write JSON report")
    args = parser.parse_args()

    duration_sec = args.minutes * 60 if args.minutes > 0 else args.hours * 3600
    end_at = time.monotonic() + duration_sec
    samples: list[dict[str, Any]] = []
    failures: list[str] = []

    pid = _backend_pid(args.port)
    if pid is None:
        print(f"ERROR: no process listening on port {args.port}", file=sys.stderr)
        return 2
    if not _health_ok(args.port):
        print(f"ERROR: backend on port {args.port} not healthy", file=sys.stderr)
        return 2

    baseline_rss = _rss_mb(pid)
    baseline_sockets = _socket_count(args.port)
    baseline_time = time.monotonic()
    print(
        f"soak: pid={pid} port={args.port} duration={duration_sec/3600:.1f}h "
        f"interval={args.interval}s baseline_rss={baseline_rss}MB sockets={baseline_sockets}\n"
    )

    while time.monotonic() < end_at:
        if not _health_ok(args.port):
            failures.append(f"health check failed at t={time.monotonic() - baseline_time:.0f}s")
            break
        pid = _backend_pid(args.port) or pid
        rss = _rss_mb(pid)
        sockets = _socket_count(args.port)
        elapsed_h = (time.monotonic() - baseline_time) / 3600.0
        sample = {
            "elapsed_sec": round(time.monotonic() - baseline_time, 1),
            "rss_mb": rss,
            "sockets": sockets,
            "health": True,
        }
        samples.append(sample)
        print(
            f"  t={sample['elapsed_sec']:6.0f}s rss={rss}MB sockets={sockets}",
            flush=True,
        )
        if baseline_rss is not None and rss is not None and elapsed_h > 0.1:
            slope = (rss - baseline_rss) / elapsed_h
            if slope > _RSS_SLOPE_MB_PER_HOUR:
                failures.append(
                    f"RSS slope {slope:.1f} MB/h exceeds SLO {_RSS_SLOPE_MB_PER_HOUR} MB/h"
                )
                break
        if sockets - baseline_sockets > _SOCKET_GROWTH_MAX:
            failures.append(
                f"socket count grew from {baseline_sockets} to {sockets} "
                f"(max growth {_SOCKET_GROWTH_MAX})"
            )
            break
        time.sleep(args.interval)

    ok = not failures
    print("\n" + "=" * 60)
    print(f"SOAK: {'PASS' if ok else 'FAIL'} ({len(samples)} samples)")
    for f in failures:
        print(f"  - {f}")
    print("=" * 60)

    report = {
        "ok": ok,
        "failures": failures,
        "baseline": {
            "pid": pid,
            "rss_mb": baseline_rss,
            "sockets": baseline_sockets,
        },
        "samples": samples,
    }
    if args.json:
        with open(args.json, "w", encoding="utf-8") as f:
            json.dump(report, f, indent=2)
        print(f"Wrote {args.json}")

    return 0 if ok else 1


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