#!/usr/bin/env python3
"""Turn the UniFi export into the VyOS config that replaces it.

Scope is deliberately narrow: DHCP and DNS. Those are the services the USG owns
that VyOS must reproduce byte-for-byte in behaviour, because getting them wrong
means clients lose their addresses or their name resolution. Everything else
either stays on UniFi (wireless), has no VyOS equivalent (user groups), or is
hand-written because it is not in the export (WAN, NAT, VRRP).

This is not a general UniFi-to-VyOS converter and should not grow into one.

  ./unifi-to-vyos.py --mode prod          # the cutover artifact
  ./unifi-to-vyos.py --mode sim           # same MACs, labsim addresses
  ./unifi-to-vyos.py --mode prod --check  # counts only, no output

Both modes come from one code path on purpose: the config proven in labsim and
the config applied to the firewalls must not be able to drift apart.

DNS note: UniFi hands out the gateway's own IP as resolver whenever a network
has no explicit dhcpd_dns -- true for 5 of the 6 VLANs, verified by labmaster
resolving against 192.168.8.1. So VyOS must run `service dns forwarding` or
those VLANs lose DNS entirely at cutover. LoT's explicit 10.0.0.194 is preserved
as-is.
"""
from __future__ import annotations

import argparse
import ipaddress
import json
import os
import re
import sys

# Where the production upstream resolver lives. Verified authoritative for
# ad.itaz.eu (SOA nas001.ad.itaz.eu) and also recursive, so it can serve as the
# single forwarder target.
UPSTREAM_DNS = "10.0.0.194"

# labsim equivalents, keyed by VLAN id. Only VLAN 10 needs a /23: every one of
# the 31 reservations is in LoT, which spans 10.0.0.x and 10.0.1.x, and a /24
# cannot represent that. k8s and Private are also /23 in production but hold no
# reservations, so they keep their existing /24 and their DHCP range is clamped
# (reported at generation time -- never silently).
SIM_SUBNETS = {
    1:   "172.31.1.0/24",
    2:   "172.31.2.0/24",
    3:   "172.31.3.0/24",
    9:   "172.31.9.0/24",
    10:  "172.31.10.0/23",
    200: "172.31.200.0/24",
}


def vlan_of(net: dict) -> int:
    """VLAN id, treating the untagged Management network as 1.

    UniFi stores vlan=None for the native network; the VyOS side already uses
    vrid 1 for it (high-availability group 'native'), so 1 is the consistent id.
    """
    return int(net["vlan"]) if net.get("vlan") else 1


def sanitize(name: str, fallback: str) -> str:
    """Reduce a UniFi client name to something VyOS will accept as a node name.

    VyOS validates static-mapping names as *hostnames*, so underscores are
    rejected outright -- verified: `Dongle-M_C0D4` fails with "Invalid static
    mapping hostname". Letters, digits and hyphens only, no leading digit or
    hyphen, no trailing hyphen.
    """
    cleaned = re.sub(r"[^A-Za-z0-9-]", "-", (name or "").strip())
    cleaned = re.sub(r"-{2,}", "-", cleaned).strip("-")
    if cleaned and cleaned[0].isdigit():
        cleaned = "h" + cleaned
    return cleaned or fallback


class Mapper:
    """Translates production addresses into the target mode's address space.

    In prod mode this is the identity. In sim mode an address is mapped by its
    offset from the network address, so the host part is preserved: 10.0.0.46
    -> 172.31.10.46 and 10.0.1.67 -> 172.31.11.67. That is what makes the sim
    test meaningful -- the MAC is identical and the host octet is recognisable.
    """

    def __init__(self, mode: str, networks: list) -> None:
        self.mode = mode
        self.clamped: list[str] = []
        self.map: dict[int, tuple] = {}
        for n in networks:
            prod = ipaddress.ip_network(
                ipaddress.ip_interface(n["subnet"]).network)
            if mode == "sim":
                sim = ipaddress.ip_network(SIM_SUBNETS[vlan_of(n)])
            else:
                sim = prod
            self.map[vlan_of(n)] = (prod, sim)

    def net(self, vlan: int) -> ipaddress.IPv4Network:
        return self.map[vlan][1]

    def addr(self, vlan: int, ip: str, what: str) -> str | None:
        """Map one address, or None if it does not fit the target subnet."""
        prod, sim = self.map[vlan]
        offset = int(ipaddress.ip_address(ip)) - int(prod.network_address)
        if offset < 0 or offset >= sim.num_addresses:
            self.clamped.append(f"{what}: {ip} does not fit {sim}")
            return None
        return str(ipaddress.ip_address(int(sim.network_address) + offset))

    def gateway(self, vlan: int, net: dict) -> str:
        """The address clients are told to use as their default route.

        Production: the USG's current address (VIPs move .254 -> .1 at cutover),
        so no client has to change anything. Sim: the sim router at .1.
        """
        if self.mode == "sim":
            return str(self.net(vlan).network_address + 1)
        return str(ipaddress.ip_interface(net["subnet"]).ip)


def build(inv: dict, mode: str) -> tuple[list[str], dict]:
    nets = [n for n in inv["networks"] if n["dhcp_enabled"] and n["subnet"]]
    nets.sort(key=vlan_of)
    m = Mapper(mode, nets)

    out: list[str] = []
    used_tags: set[str] = set()
    stats = {"subnets": 0, "mappings": 0, "dropped": []}
    by_vlan: dict[int, list] = {}
    for r in inv["reservations"]:
        if r["network_vlan"] is None and r["network_name"] != "Management":
            # resolved_by == "unresolved"; cannot place it without a subnet
            stats["dropped"].append(f"{r['ip']} {r['mac']} (no network)")
            continue
        by_vlan.setdefault(r["network_vlan"] or 1, []).append(r)

    out.append("# --- DHCP ---------------------------------------------------")
    for n in nets:
        vlan = vlan_of(n)
        sub = m.net(vlan)
        base = f"set service dhcp-server shared-network-name {sanitize(n['name'], f'vlan{vlan}')} subnet {sub}"
        gw = m.gateway(vlan, n)

        out.append("")
        out.append(f"# {n['name']} (VLAN {vlan}) <- {n['subnet']}")
        # subnet-id is required by kea and must be stable across regenerations;
        # the VLAN id is already the unique per-network number in this lab.
        out.append(f"{base} subnet-id {vlan}")
        out.append(f"{base} option default-router {gw}")

        # Preserve UniFi's explicit resolver where it set one (LoT -> the NAS);
        # otherwise hand out the gateway, which is what the USG does today and
        # what `service dns forwarding` below will answer on.
        for ns in (n["dhcp_dns"] or [gw]):
            mapped = ns if ns not in (UPSTREAM_DNS,) or mode == "prod" else gw
            if mode == "sim" and ns == UPSTREAM_DNS:
                mapped = gw  # no NAS in the sim; the router resolves
            out.append(f"{base} option name-server {mapped}")

        if n["domain_name"]:
            out.append(f"{base} option domain-name '{n['domain_name']}'")
        if n["dhcp_lease"]:
            out.append(f"{base} lease {n['dhcp_lease']}")

        start = m.addr(vlan, n["dhcp_start"], f"{n['name']} range start")
        stop = m.addr(vlan, n["dhcp_stop"], f"{n['name']} range stop")
        if start is None:
            start = str(sub.network_address + 11)
        if stop is None:
            # Clamp to the last usable address rather than dropping the pool.
            stop = str(sub.broadcast_address - 1)
        out.append(f"{base} range LAN start {start}")
        out.append(f"{base} range LAN stop {stop}")
        stats["subnets"] += 1

        for r in sorted(by_vlan.get(vlan, []), key=lambda x: ipaddress.ip_address(x["ip"])):
            ip = m.addr(vlan, r["ip"], f"reservation {r['name']}")
            if ip is None:
                stats["dropped"].append(f"{r['ip']} {r['mac']} ({r['name']})")
                continue
            tag = sanitize(r["name"] or r["hostname"], "host-" + r["mac"].replace(":", ""))
            # Distinct clients can sanitize to the same name; a collision would
            # silently overwrite one reservation with another's address.
            if tag in used_tags:
                tag = f"{tag}-{r['mac'].replace(':', '')[-4:]}"
            used_tags.add(tag)
            out.append(f"{base} static-mapping {tag} mac {r['mac']}")
            out.append(f"{base} static-mapping {tag} ip-address {ip}")
            stats["mappings"] += 1

    out.append("")
    out.append("# --- DNS ----------------------------------------------------")
    out.append("# The USG resolves for 5 of 6 VLANs today (it hands out its own")
    out.append("# address when dhcpd_dns is empty). Without this, they lose DNS.")
    for n in nets:
        vlan = vlan_of(n)
        out.append(f"set service dns forwarding listen-address {m.gateway(vlan, n)}")
        out.append(f"set service dns forwarding allow-from {m.net(vlan)}")
    if mode == "prod":
        out.append(f"set service dns forwarding name-server {UPSTREAM_DNS}")
    else:
        # The sim has no NAS; forward to whatever the sim host can reach.
        out.append("set service dns forwarding name-server 1.1.1.1")
    out.append("set service dns forwarding cache-size 10000")

    stats["clamped"] = m.clamped
    return out, stats


def main() -> int:
    ap = argparse.ArgumentParser()
    here = os.path.dirname(os.path.abspath(__file__))
    ap.add_argument("--mode", choices=("prod", "sim"), required=True)
    ap.add_argument("--inventory", default=os.path.join(here, "export", "inventory.json"))
    ap.add_argument("-o", "--out")
    ap.add_argument("--check", action="store_true", help="counts only, no config")
    args = ap.parse_args()

    with open(args.inventory) as fh:
        inv = json.load(fh)

    lines, stats = build(inv, args.mode)

    expected = len(inv["reservations"])
    print(f"mode={args.mode}  subnets={stats['subnets']}  "
          f"static-mappings={stats['mappings']}/{expected}", file=sys.stderr)
    for c in stats["clamped"]:
        print(f"  clamped: {c}", file=sys.stderr)
    for d in stats["dropped"]:
        print(f"  DROPPED: {d}", file=sys.stderr)

    if stats["mappings"] != expected and args.mode == "prod":
        print(f"ERROR: {expected - stats['mappings']} reservation(s) missing from "
              f"prod output -- every one must survive the cutover", file=sys.stderr)
        return 1

    if args.check:
        return 0

    text = "\n".join(lines) + "\n"
    if args.out:
        with open(args.out, "w") as fh:
            fh.write(text)
        print(f"wrote {args.out}", file=sys.stderr)
    else:
        sys.stdout.write(text)
    return 0


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