#!/usr/bin/env python3
"""Set a fixed-IP reservation in UniFi, so it cannot lease that address away.

Written for the firewalls' management NICs: their addresses are pinned static
on the VyOS side, but UniFi still owns the pool they sit in and would happily
hand the same address to something else. A reservation closes that gap while
the USG is still the DHCP server, and it keeps the management address identical
in both cutover modes.

  ./unifi-reserve.py 64:62:66:25:96:47 192.168.8.143
  ./unifi-reserve.py --dry-run <mac> <ip>

Idempotent: an existing, matching reservation is reported and left alone. This
writes to the live controller, so it verifies by reading the record back rather
than trusting the response.
"""
from __future__ import annotations

import argparse
import ipaddress
import sys

import _unifi


def find_network(nets: list, ip: str) -> dict | None:
    """Which configured network contains this address?

    ip_subnet holds the gateway address with a prefix ("192.168.8.1/23"), so
    the network has to be derived from it rather than compared directly.
    """
    addr = ipaddress.ip_address(ip)
    for n in nets:
        raw = n.get("ip_subnet")
        if not raw:
            continue
        try:
            if addr in ipaddress.ip_interface(raw).network:
                return n
        except ValueError:
            continue
    return None


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("mac")
    ap.add_argument("ip")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    mac = args.mac.lower().replace("-", ":")
    ipaddress.ip_address(args.ip)  # fail early on a typo

    opener, base, site = _unifi.client()
    users = _unifi.get(opener, base, site, "rest/user")
    nets = _unifi.get(opener, base, site, "rest/networkconf")
    for blob in (users, nets):
        if isinstance(blob, dict):
            print(f"error reading controller: {blob['__error__']}", file=sys.stderr)
            return 1

    user = next((u for u in users if (u.get("mac") or "").lower() == mac), None)
    if user is None:
        print(f"{mac} is not a known client -- connect it once, or create it "
              f"in the UI first", file=sys.stderr)
        return 1

    net = find_network(nets, args.ip)
    if net is None:
        print(f"no configured network contains {args.ip}", file=sys.stderr)
        return 1

    label = user.get("name") or user.get("hostname") or mac
    if user.get("use_fixedip") and user.get("fixed_ip") == args.ip:
        print(f"{label} ({mac}) already reserved at {args.ip} -- nothing to do")
        return 0
    if user.get("use_fixedip"):
        print(f"WARNING: {label} currently reserved at {user.get('fixed_ip')}, "
              f"changing to {args.ip}", file=sys.stderr)

    print(f"{label} ({mac}) -> {args.ip} on '{net.get('name')}' (VLAN {net.get('vlan') or 'native'})")
    if args.dry_run:
        print("  --dry-run: not writing")
        return 0

    payload = {"use_fixedip": True, "fixed_ip": args.ip, "network_id": net["_id"]}
    res = _unifi.put(opener, base, site, f"rest/user/{user['_id']}", payload)
    if isinstance(res, dict):
        print(f"  write failed: {res['__error__']}", file=sys.stderr)
        return 1

    # Read it back: the controller accepting a PUT is not proof it stored what
    # we asked for.
    after = _unifi.get(opener, base, site, "rest/user")
    check = next((u for u in after if (u.get("mac") or "").lower() == mac), {})
    if check.get("use_fixedip") and check.get("fixed_ip") == args.ip:
        print(f"  verified: reservation is live")
        return 0
    print(f"  VERIFY FAILED: controller reports use_fixedip="
          f"{check.get('use_fixedip')} fixed_ip={check.get('fixed_ip')}", file=sys.stderr)
    return 1


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