#!/usr/bin/env python3
"""Probe backend MLX auto-restart resilience.

Sends a lightweight chat probe, optionally signals MLX stop (if the prefs route
is reachable), then re-probes until recovery or timeout.

Usage:
  ~/.mugi/venv/bin/python scripts/mlx_restart_resilience.py
  ~/.mugi/venv/bin/python scripts/mlx_restart_resilience.py --skip-stop

Requires auth token in MUGI_AUTH_TOKEN or discovers backend like live batteries.
"""
from __future__ import annotations

import argparse
import os
import sys
import time

_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _REPO not in sys.path:
    sys.path.insert(0, _REPO)

from utils.mcp_backend_client import discover_backend, http_json, _headers  # noqa: E402


def _mlx_status(conn) -> dict:
    code, out = http_json(
        "GET",
        f"{conn.base_url}/preferences/mlx/status",
        _headers(conn),
        timeout=15.0,
    )
    return out if code == 200 and isinstance(out, dict) else {"error": out, "code": code}


def main() -> int:
    parser = argparse.ArgumentParser(description="MLX restart resilience probe")
    parser.add_argument("--skip-stop", action="store_true", help="Only poll status (no stop)")
    parser.add_argument("--timeout", type=float, default=120.0)
    parser.add_argument("--poll", type=float, default=5.0)
    args = parser.parse_args()

    conn = discover_backend()
    if isinstance(conn, dict):
        print("BACKEND:", conn, file=sys.stderr)
        return 2

    before = _mlx_status(conn)
    print(f"mlx status before: {before}")

    if not args.skip_stop:
        code, out = http_json(
            "POST",
            f"{conn.base_url}/preferences/mlx/stop",
            _headers(conn),
            {},
            timeout=30.0,
        )
        print(f"mlx stop: {code} {out}")

    deadline = time.monotonic() + args.timeout
    recovered = False
    while time.monotonic() < deadline:
        st = _mlx_status(conn)
        running = bool(st.get("running") or st.get("server_running"))
        print(f"  poll: running={running} detail={st!r:.200}")
        if running:
            recovered = True
            break
        time.sleep(args.poll)

    if recovered:
        print("MLX RESTART RESILIENCE: PASS (server running again)")
        return 0
    print("MLX RESTART RESILIENCE: FAIL (did not recover within timeout)")
    return 1


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