#!/usr/bin/env python3
"""Reboot the HTPC as soon as an SSH window opens, then confirm it came back.

Proof of reboot is a changed kernel boot_id, not the SSH session dropping --
on this link a dropped session is the normal case and proves nothing.
"""
import os, subprocess, sys, time
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from suvi import gpubox

def ssh(command, stdin=None, timeout=40):
    try:
        r = subprocess.run(gpubox.ssh_argv() + [command], input=stdin,
                           stdout=subprocess.PIPE, stderr=subprocess.STDOUT, timeout=timeout)
        return r.returncode, r.stdout.decode(errors="replace")
    except subprocess.TimeoutExpired:
        return -1, "timeout"

def log(msg):
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)

# 1. Wait for a window, capture boot_id and issue the reboot in the same round trip.
#    Polkit usually refuses reboot from an SSH session while a desktop session is
#    open, so fall back to sudo with the password on stdin (never on the command line).
before = None
for attempt in range(1, 361):
    code, out = ssh("echo BOOT=$(cat /proc/sys/kernel/random/boot_id); "
                    "systemctl reboot --no-ask-password 2>/dev/null "
                    "|| sudo -S -p '' systemctl reboot; echo RC=$?",
                    stdin=(gpubox._password() + "\n").encode())
    for line in out.splitlines():
        if line.startswith("BOOT="):
            before = line[5:].strip()
    if before:
        log(f"attempt {attempt}: connected, boot_id {before[:8]}..., reboot issued "
            f"({[l for l in out.splitlines() if l.startswith('RC=')] or 'session dropped'})")
        break
    time.sleep(20)
else:
    log("no SSH window in ~2 h; reboot NOT issued"); sys.exit(1)

# 2. Wait for it to come back with a different boot_id.
time.sleep(60)
for attempt in range(1, 361):
    code, out = ssh("echo BOOT=$(cat /proc/sys/kernel/random/boot_id); uptime -p; "
                    "echo LLAMA=$(systemctl --user is-active llama-swap.service); "
                    "iw dev wlan0 link 2>/dev/null | grep -E 'signal|bitrate'")
    now = next((l[5:].strip() for l in out.splitlines() if l.startswith("BOOT=")), None)
    if now and now != before:
        log(f"back after reboot: boot_id {now[:8]}... (was {before[:8]}...)")
        for line in out.splitlines():
            if not line.startswith("BOOT="):
                log("  " + line.strip())
        break
    if now == before:
        log("still the old boot -- reboot has not happened yet")
    time.sleep(20)
else:
    log("did not come back within ~2 h"); sys.exit(1)

# 3. Link quality now, the point of the exercise.
p = subprocess.run(["ping", "-c", "30", "-i", "1", "-W", "2", gpubox.HOST.split("@")[-1]],
                   stdout=subprocess.PIPE).stdout.decode()
log("ping: " + " / ".join(l.strip() for l in p.splitlines()[-2:]))