#!/usr/bin/env python3
#
# ap-status: report connectivity and WiFi client status for managed APs.
#
# Output is JSON to stdout in the format:
#   {"aps": [{"id": ..., "ip": ..., "mac": ..., "ok": ..., "status": ..., "clients": [...]}]}
#
# Status values (ok=true):
#   healthy        - ping, SSH, and station dump all succeeded
#
# Status values (ok=false):
#   config_error   - ping OK but SSH access failed (key mismatch, AP reset, etc.)
#   config_missing - AP responds to ping but is not in ap_manager state (full-range scan only)
#   ap_error       - SSH OK but station dump or other AP-side command failed

import os
import sys
import json
import subprocess

STATE_PATH = "/var/lib/ap_manager/state"
SSH_KEY = f"{STATE_PATH}/.ssh/id_apmanager"
SSH_OPTS = [
    "-i", SSH_KEY,
    "-o", "StrictHostKeyChecking=no",
    "-o", "UserKnownHostsFile=/dev/null",
    "-o", "ConnectTimeout=5",
    "-o", "BatchMode=yes",
]
SUBNET_START = 11
MAX_NUMBER = 20
PING_TIMEOUT = 2

HELP_MSG = """
Usage: ap-status [OPTIONS]

Options:
  -h, --help  Print usage and exit
  -n INDEX    Check a single AP by index (1-based)
  -a          Check full IP range (*.11-*.30); APs not in state get status
              "config missing"

By default only configured APs (from ap_manager state) are checked.
APs that do not respond to ping are silently skipped.
"""

ap_index = None
all_range = False

args = sys.argv[1:]
i = 0
while i < len(args):
    arg = args[i]
    if arg in ('-h', '--help', 'help'):
        print(HELP_MSG)
        sys.exit(0)
    elif arg == '-n':
        i += 1
        if i >= len(args):
            print("ERROR: -n requires an argument")
            sys.exit(1)
        try:
            ap_index = int(args[i])
        except ValueError:
            print(f"ERROR: invalid index '{args[i]}'")
            sys.exit(1)
        if not (1 <= ap_index <= MAX_NUMBER):
            print(f"ERROR: index must be between 1 and {MAX_NUMBER}")
            sys.exit(1)
    elif arg == '-a':
        all_range = True
    else:
        print(f'Invalid argument "{arg}"')
        sys.exit(1)
    i += 1

if ap_index is not None and all_range:
    print("ERROR: -n and -a are mutually exclusive")
    sys.exit(1)


def ping_host(ip):
    r = subprocess.run(
        ["ping", "-W", str(PING_TIMEOUT), "-c", "1", ip],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
    )
    return r.returncode == 0


def get_mac(ip):
    r = subprocess.run(["ip", "neigh", "show", ip], capture_output=True, text=True)
    for part in r.stdout.split():
        if part == "lladdr":
            idx = r.stdout.split().index(part)
            return r.stdout.split()[idx + 1]
    return None


def ssh_run(ip, cmd, timeout=10):
    try:
        r = subprocess.run(
            ["ssh", "-q"] + SSH_OPTS + [f"root@{ip}", cmd],
            capture_output=True, text=True, timeout=timeout,
        )
        return r.returncode, r.stdout
    except subprocess.TimeoutExpired:
        return 1, ""


def parse_station_dump(output):
    clients = []
    current = None
    for line in output.splitlines():
        s = line.strip()
        if s.startswith("Station "):
            if current:
                clients.append(current)
            parts = s.split()
            current = {"mac": parts[1], "rssi": 0, "connected_s": 0}
        elif current is None:
            continue
        elif s.startswith("signal:"):
            # value may be "-65 dBm" or "-65 [-65, -66] dBm" (multi-chain)
            val = s.split(":", 1)[1].strip().split()[0]
            try:
                current["rssi"] = int(val)
            except ValueError:
                pass
        elif s.startswith("connected time:"):
            val = s.split(":", 1)[1].strip().split()[0]
            try:
                current["connected_s"] = int(val)
            except ValueError:
                pass
    if current:
        clients.append(current)
    return clients


def check_ap(idx, prefix, configured=True):
    ip = f"{prefix}.{SUBNET_START + idx - 1}"
    ap_id = idx

    if not ping_host(ip):
        return None

    mac = get_mac(ip)

    if not configured:
        return {"id": ap_id, "ip": ip, "mac": mac, "ok": False, "status": "config_missing", "clients": []}

    rc, out = ssh_run(ip, "echo ok")
    if rc != 0 or out.strip() != "ok":
        return {"id": ap_id, "ip": ip, "mac": mac, "ok": False, "status": "config_error", "clients": []}

    rc, out = ssh_run(ip, "iw dev | awk '/Interface/{iface=$2} /type AP/{print iface}'")
    if rc != 0:
        return {"id": ap_id, "ip": ip, "mac": mac, "ok": False, "status": "ap_error", "clients": []}

    interfaces = [l.strip() for l in out.splitlines() if l.strip()]
    if not interfaces:
        return {"id": ap_id, "ip": ip, "mac": mac, "ok": False, "status": "ap_error", "clients": []}

    clients = []
    for iface in interfaces:
        rc, dump = ssh_run(ip, f"iw dev {iface} station dump")
        if rc != 0:
            return {"id": ap_id, "ip": ip, "mac": mac, "ok": False, "status": "ap_error", "clients": []}
        clients.extend(parse_station_dump(dump))

    return {"id": ap_id, "ip": ip, "mac": mac, "ok": True, "status": "healthy", "clients": clients}


def load_state():
    state_file = f"{STATE_PATH}/current"
    next_number = 0
    network = "192.168.14.0"
    if os.path.exists(state_file):
        with open(state_file) as f:
            for line in f:
                line = line.strip()
                if line.startswith("next_number="):
                    try:
                        next_number = int(line.split("=", 1)[1])
                    except ValueError:
                        pass
                elif line.startswith("network="):
                    network = line.split("=", 1)[1].strip('"')
    return next_number, network


next_number, network = load_state()
prefix = network.rsplit(".", 1)[0]

if ap_index is not None:
    indices = [ap_index]
elif all_range:
    indices = range(1, MAX_NUMBER + 1)
else:
    indices = range(1, next_number)

aps = []
for idx in indices:
    is_configured = not all_range or idx < next_number
    result = check_ap(idx, prefix, configured=is_configured)
    if result is not None:
        aps.append(result)

print(json.dumps({"aps": aps}))