#!/usr/bin/env python3
"""Change which access points broadcast a WiFi network (SSID) on one UniFi console.

Usage:
    docs/examples/waters-church-na/wifi-broadcast-devices.py north-attleboro "Waters Guest" --list
    docs/examples/waters-church-na/wifi-broadcast-devices.py north-attleboro "Waters Guest" \
        --add "Auditorium Left,Auditorium Right,Main Lobby Entrance"           # dry run
    docs/examples/waters-church-na/wifi-broadcast-devices.py north-attleboro "Waters Guest" \
        --add "Auditorium Left,Auditorium Right,Main Lobby Entrance" --apply   # write it
    docs/examples/waters-church-na/wifi-broadcast-devices.py north-attleboro "Waters Guest" --remove "The Vault" --apply
    docs/examples/waters-church-na/wifi-broadcast-devices.py north-attleboro "Waters Guest" --all --apply   # every AP

Names are matched exactly against the console's device names (--list prints
them). Without --apply nothing is written: the script prints the APs the SSID
is on now and the list it would send. The whole broadcast object is PUT back
with only broadcastingDeviceFilter changed.

Needs the per-console Network API key at ~/.config/unifi/<site>-network-api-key
(see unifi-dns-record.py). The key is only ever sent in the X-API-KEY header.
"""

from __future__ import annotations

import argparse
import json
import os
import ssl
import sys
import urllib.error
import urllib.request

SITES = {
    "apollo-beach": "10.3.0.1",
    "north-attleboro": "10.1.0.1",
    "woonsocket": "10.2.0.1",
    "norwood": "10.4.0.1",
    "southcoast": "10.5.0.1",
}
KEY_DIR = os.path.expanduser("~/.config/unifi")
TIMEOUT = 20


class Console:
    def __init__(self, name: str):
        self.name = name
        self.gateway = SITES[name]
        with open(os.path.join(KEY_DIR, f"{name}-network-api-key")) as f:
            self.key = f.read().strip()
        self.ctx = ssl.create_default_context()
        self.ctx.check_hostname = False
        self.ctx.verify_mode = ssl.CERT_NONE

    def request(self, method: str, path: str, body: dict | None = None):
        url = f"https://{self.gateway}/proxy/network/integration/v1{path}"
        data = json.dumps(body).encode() if body is not None else None
        req = urllib.request.Request(url, data=data, method=method, headers={
            "X-API-KEY": self.key, "Accept": "application/json",
            "Content-Type": "application/json"})
        try:
            with urllib.request.urlopen(req, timeout=TIMEOUT, context=self.ctx) as r:
                raw = r.read().decode()
        except urllib.error.HTTPError as e:
            raise SystemExit(f"{method} {path} -> HTTP {e.code}: {e.read().decode(errors='replace')[:600]}")
        return json.loads(raw) if raw else None

    def site_id(self) -> str:
        return self.request("GET", "/sites")["data"][0]["id"]

    def access_points(self, sid: str) -> dict[str, str]:
        devs = self.request("GET", f"/sites/{sid}/devices?limit=200")["data"]
        return {d["name"]: d["id"] for d in devs if "accessPoint" in (d.get("features") or [])}

    def broadcast(self, sid: str, ssid: str) -> dict:
        for b in self.request("GET", f"/sites/{sid}/wifi/broadcasts?limit=200")["data"]:
            if b.get("name") == ssid:
                return b
        raise SystemExit(f"no WiFi network named {ssid!r} on {self.name}")


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("site", choices=sorted(SITES))
    ap.add_argument("ssid", help="WiFi network name as shown in the console")
    ap.add_argument("--add", default="", help="comma-separated AP names to add")
    ap.add_argument("--remove", default="", help="comma-separated AP names to remove")
    ap.add_argument("--all", action="store_true", help="broadcast on every AP (clears the device filter)")
    ap.add_argument("--list", action="store_true", help="print the APs and where the SSID broadcasts, then exit")
    ap.add_argument("--apply", action="store_true", help="actually write the change")
    a = ap.parse_args()

    c = Console(a.site)
    sid = c.site_id()
    aps = c.access_points(sid)
    by_id = {v: k for k, v in aps.items()}
    b = c.broadcast(sid, a.ssid)
    filt = b.get("broadcastingDeviceFilter") or {}
    current = list(filt.get("deviceIds", [])) if filt.get("type") == "DEVICES" else None

    def names(ids):
        return ", ".join(by_id.get(i, i) for i in ids)

    print(f"{a.site}: {a.ssid!r} broadcasts on " + ("ALL access points" if current is None else f"{len(current)} AP(s): {names(current)}"))
    if a.list:
        print("\nAccess points on this console:")
        for n in sorted(aps):
            print(f"  {'*' if current is None or aps[n] in current else ' '} {n}")
        return 0

    add = [s.strip() for s in a.add.split(",") if s.strip()]
    rem = [s.strip() for s in a.remove.split(",") if s.strip()]
    unknown = [n for n in add + rem if n not in aps]
    if unknown:
        raise SystemExit(f"unknown AP name(s): {', '.join(unknown)}; use --list")
    if not (add or rem or a.all):
        ap.error("nothing to change: give --add, --remove or --all")

    new_b = {k: v for k, v in b.items() if k not in ("id", "metadata")}
    if a.all:
        new_b.pop("broadcastingDeviceFilter", None)
        summary = "ALL access points"
    else:
        ids = list(current or [])
        for n in add:
            if aps[n] not in ids:
                ids.append(aps[n])
        ids = [i for i in ids if i not in {aps[n] for n in rem}]
        new_b["broadcastingDeviceFilter"] = {"type": "DEVICES", "deviceIds": ids}
        summary = f"{len(ids)} AP(s): {names(ids)}"
    print(f"would broadcast on {summary}")
    if not a.apply:
        print("DRY RUN, add --apply to write it")
        return 0
    out = c.request("PUT", f"/sites/{sid}/wifi/broadcasts/{b['id']}", new_b)
    f2 = (out or {}).get("broadcastingDeviceFilter") or {}
    print("UPDATED: now on " + ("ALL access points" if f2.get("type") != "DEVICES" else names(f2.get("deviceIds", []))))
    return 0


if __name__ == "__main__":
    sys.exit(main())