#!/usr/bin/env python3
import subprocess
import time

INTERFACE = "wg0"
RECENT_HANDSHAKE_SECONDS = 30 * 60


def _failure(detail):
    return {
        "state": "offline", "health": "unavailable", "summary": "Unavailable",
        "detail": detail, "interface": INTERFACE, "interface_active": False,
        "configured_peers": 0, "recent_handshakes": 0, "stale_handshakes": 0,
        "rx_bytes": 0, "tx_bytes": 0,
        "recent_handshake_threshold_seconds": RECENT_HANDSHAKE_SECONDS,
    }


def _format_bytes(value):
    units = ("B", "KiB", "MiB", "GiB", "TiB")
    amount = float(value)
    for unit in units:
        if amount < 1024 or unit == units[-1]:
            return f"{amount:.2f} {unit}" if unit != "B" else f"{int(amount)} B"
        amount /= 1024


def collect():
    result = subprocess.run(["wg", "show", INTERFACE, "dump"],
                            capture_output=True, text=True)
    if result.returncode != 0:
        return _failure("WireGuard telemetry unavailable")
    lines = [line.split() for line in result.stdout.splitlines() if line.strip()]
    if not lines or len(lines[0]) < 4:
        return _failure("WireGuard returned invalid telemetry")
    peers = lines[1:]
    now = int(time.time())
    recent = rx_bytes = tx_bytes = 0
    try:
        for peer in peers:
            if len(peer) < 8:
                return _failure("WireGuard returned incomplete telemetry")
            handshake = int(peer[4])
            rx_bytes += int(peer[5])
            tx_bytes += int(peer[6])
            if handshake and now - handshake <= RECENT_HANDSHAKE_SECONDS:
                recent += 1
    except (TypeError, ValueError):
        return _failure("WireGuard returned invalid telemetry")
    configured = len(peers)
    stale = configured - recent
    state = "online" if configured and recent == configured else "warning"
    health = "healthy" if state == "online" else "partial"
    return {
        "state": state, "health": health,
        "summary": f"{recent}/{configured} peers recent",
        "detail": (f"WireGuard active · {_format_bytes(rx_bytes)} RX · "
                   f"{_format_bytes(tx_bytes)} TX"),
        "interface": INTERFACE, "interface_active": True,
        "configured_peers": configured, "recent_handshakes": recent,
        "stale_handshakes": stale, "rx_bytes": rx_bytes, "tx_bytes": tx_bytes,
        "recent_handshake_threshold_seconds": RECENT_HANDSHAKE_SECONDS,
    }