#!/usr/bin/env python3
"""Bound a busy llama.cpp slot even if the client loses its cancellation path.

Only task IDs, monotonic timestamps and service invocation IDs are persisted.
The wrapper holds the same maintenance lock as the updater and restart helper.
"""
import json
import os
from pathlib import Path
import subprocess
import time
from urllib.request import urlopen


def observe(previous, slots, invocation, now, seconds, tokens):
    old = previous if previous.get("invocation") == invocation else {}
    active = {}
    reasons = []
    for slot in slots:
        if not slot.get("is_processing"):
            continue
        if "id_task" not in slot:
            raise ValueError("busy slot without task ID")
        key = f'{slot["id"]}:{slot["id_task"]}'
        started = old.get("active", {}).get(key, now)
        if started > now:
            started = now
        active[key] = started
        decoded = max((item.get("n_decoded", 0) for item in slot.get("next_token", [])), default=0)
        if now - started >= seconds:
            reasons.append(f"slot {slot['id']}: request exceeded {seconds}s")
        if decoded > tokens:
            reasons.append(f"slot {slot['id']}: output exceeded {tokens} tokens")
    return {"invocation": invocation, "active": active}, reasons


def systemctl(*args):
    return subprocess.run(["systemctl", "--user", *args], check=True,
                          capture_output=True, text=True, timeout=35).stdout.strip()


def main():
    service = os.environ.get("LLAMA_SERVICE", "llama-qwen.service")
    url = os.environ.get("LLAMA_URL", "http://127.0.0.1:8080").rstrip("/")
    seconds = int(os.environ.get("LLAMA_MAX_REQUEST_SECONDS", "600"))
    tokens = int(os.environ.get("LLAMA_MAX_OUTPUT_TOKENS", "2048"))
    if seconds <= 0 or tokens <= 0:
        raise ValueError("watchdog limits must be positive")
    state_dir = Path(os.environ.get("XDG_STATE_HOME", str(Path.home() / ".local/state"))) / "hermes-stack-maintenance"
    state_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
    path = state_dir / "llama-watchdog.json"
    try:
        previous = json.loads(path.read_text())
    except (OSError, ValueError):
        previous = {}
    boot = Path("/proc/sys/kernel/random/boot_id").read_text().strip()
    if previous.get("boot") != boot:
        previous = {}
    status = systemctl("show", service, "--property=ActiveState", "--property=InvocationID")
    info = dict(line.split("=", 1) for line in status.splitlines())
    if info.get("ActiveState") != "active":
        return  # Respect explicit stops and maintenance; never start an inactive server.
    invocation = info["InvocationID"]
    now = time.monotonic()
    try:
        with urlopen(url + "/slots", timeout=5) as response:
            slots = json.load(response)
        if not isinstance(slots, list):
            raise ValueError("invalid slot response")
        state, reasons = observe(previous, slots, invocation, now, seconds, tokens)
    except (OSError, ValueError, KeyError, TypeError):
        # A single failed probe is not evidence of a hung model. Require a full
        # request budget of consecutive failures within the same server process.
        since = previous.get("unreachable_since", now) if previous.get("invocation") == invocation else now
        state = {"invocation": invocation, "active": previous.get("active", {}), "unreachable_since": since}
        reasons = [f"slot endpoint unreachable for {seconds}s"] if now - since >= seconds else []
    last_restart = previous.get("last_restart", -1e12)
    state["last_restart"] = last_restart
    if reasons and now - last_restart >= max(300, seconds):
        print("llama watchdog: " + "; ".join(reasons) + "; restarting " + service, flush=True)
        systemctl("restart", "--no-block", service)
        state = {"invocation": invocation, "active": {}, "last_restart": now}
    state["boot"] = boot
    pending = path.with_suffix(".pending")
    pending.write_text(json.dumps(state) + "\n")
    pending.chmod(0o600)
    pending.replace(path)


if __name__ == "__main__":
    main()