#!/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

# Upstream resolvers for VyOS's own forwarder -- the same pair the USG used on
# its WAN (wan_dns1/wan_dns2). The NAS at 10.0.0.194 is deliberately NOT here:
# it is legacy for ad.itaz.eu, and those records now live in Cloudflare, so the
# zone resolves publicly (verified: nas001.ad.itaz.eu and kvm-macstudio1 both
# answer from 8.8.8.8). That means no conditional forward is needed and the NAS
# is out of the DNS path entirely.
UPSTREAM_DNS = ["8.8.8.8", "8.8.4.4"]

# 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}")

        # Every VLAN is handed the gateway as its resolver, so all lookups go
        # through VyOS and out to the upstreams above. UniFi set an explicit
        # resolver on LoT only (the NAS); that is deliberately not carried over
        # -- the NAS is legacy and pointing clients at it would keep it in the
        # path for one VLAN and not the others.
        out.append(f"{base} option name-server {gw}")

        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)}")
    for ns in UPSTREAM_DNS:
        out.append(f"set service dns forwarding name-server {ns}")
    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())