// ipadm - dhcp/dns management for the dnsmasq-based router setup on this host.
//
// Maintains a flat-file database of static hosts (hostname, IP, optional
// MAC, comment) and of port-forwards (WAN port/proto -> host from that
// database). On `ipadm -u`, regenerates the dnsmasq include file
// (host-record/dhcp-host entries) and the nftables include file (DNAT
// rules), then reloads both services.
package main

import (
	"fmt"
	"os"
	"strings"
	"text/tabwriter"
)

const version = "1.2.0"

type usageRow struct {
	left string
	desc string
	wrap bool // put desc on its own line, indented to the description column
}

var usageRows = []usageRow{
	{"ipadm <hostname>", "add/edit host (interactive)", false},
	{"ipadm -l", "list hosts", false},
	{"ipadm -a [-f] <hostname> [ip address] [mac address]", "add host (IP auto-assigned if omitted)", false},
	{"ipadm -c <hostname> <comment>", "set comment", false},
	{"ipadm -i <hostname> <ip address>", "change ip address", false},
	{"ipadm -m <hostname> <mac address>", "change mac address", false},
	{"ipadm -r <old hostname> <new hostname>", "rename host", false},
	{"ipadm -d <hostname>", "delete host", false},
	{"ipadm -pa [-f] <hostname> <wan-port> [lan-port] [tcp|udp|both]", "add port-forward to a known host", true},
	{"ipadm -pl", "list port-forwards", false},
	{"ipadm -pd <wan-port> [tcp|udp|both]", "delete port-forward", false},
	{"ipadm -u", "update (regenerate dnsmasq+nftables config, apply)", false},
	{"ipadm -h", "this help", false},
}

// buildUsage lays out usageRows in two aligned columns. Rows too long to
// share a line with their description (marked wrap) get the description on
// the next line instead, so one oversized entry can't drag every other
// line's alignment along with it.
func buildUsage() string {
	const gap = 3
	width := 0
	for _, r := range usageRows {
		if !r.wrap && len(r.left) > width {
			width = len(r.left)
		}
	}

	var b strings.Builder
	fmt.Fprintf(&b, "ipadm - dhcp/dns management (dnsmasq backend) v%s\n\n", version)
	for i, r := range usageRows {
		prefix := "       "
		if i == 0 {
			prefix = "usage: "
		}
		if r.wrap {
			fmt.Fprintf(&b, "%s%s\n", prefix, r.left)
			fmt.Fprintf(&b, "%s%s%s\n", prefix, strings.Repeat(" ", width+gap), r.desc)
		} else {
			fmt.Fprintf(&b, "%s%-*s%s%s\n", prefix, width, r.left, strings.Repeat(" ", gap), r.desc)
		}
	}
	return b.String()
}

var usage = buildUsage()

// netConfig describes the LAN/WAN the tool validates against. Overridable
// via env vars, mainly so this can be tested without touching the real
// /etc files.
//
// The LAN is a /22 split into four /24s: staticCIDR for hosts without a MAC
// (manually configured on the device, DNS only), reservedCIDR for hosts with
// a MAC (DHCP static reservation), plus a dynamic DHCP pool and a reserved,
// currently unused /24 that ipadm doesn't need to know about.
type netConfig struct {
	staticCIDR   string
	reservedCIDR string
	gatewayIP    string
	wanIface     string
}

// cidrFor returns the subnet a host with the given MAC (possibly empty)
// belongs in.
func (c netConfig) cidrFor(mac string) string {
	if mac == "" {
		return c.staticCIDR
	}
	return c.reservedCIDR
}

func envOr(key, def string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return def
}

func dbPath() string {
	return envOr("IPADM_DB", "/etc/ipadm/hosts")
}

func portFwdDBPath() string {
	return envOr("IPADM_PORTFWD_DB", "/etc/ipadm/portforwards")
}

func dnsmasqHostsPath() string {
	return envOr("IPADM_DNSMASQ_HOSTS", "/etc/dnsmasq.d/hosts.conf")
}

func nftPortFwdPath() string {
	return envOr("IPADM_NFT_PORTFWD", "/etc/nftables.d/portforward.conf")
}

func loadNetConfig() netConfig {
	return netConfig{
		staticCIDR:   envOr("IPADM_STATIC_CIDR", "10.0.0.0/24"),
		reservedCIDR: envOr("IPADM_RESERVED_CIDR", "10.0.1.0/24"),
		gatewayIP:    envOr("IPADM_GATEWAY_IP", "10.0.0.1"),
		wanIface:     envOr("IPADM_WAN_IFACE", "enp3s0"),
	}
}

func fail(format string, a ...any) {
	fmt.Fprintf(os.Stderr, "ipadm: "+format+"\n", a...)
	os.Exit(1)
}

func main() {
	args := os.Args[1:]
	if len(args) == 0 {
		fmt.Fprint(os.Stderr, usage)
		os.Exit(1)
	}

	switch args[0] {
	case "-h", "--help":
		fmt.Print(usage)
	case "-l":
		requireArgs(args[1:], 0, "ipadm -l")
		cmdList()
	case "-u":
		requireArgs(args[1:], 0, "ipadm -u")
		cmdUpdate()
	case "-a":
		cmdAdd(args[1:])
	case "-c":
		cmdComment(args[1:])
	case "-i":
		cmdChangeIP(args[1:])
	case "-m":
		cmdChangeMAC(args[1:])
	case "-r":
		cmdRename(args[1:])
	case "-d":
		cmdDelete(args[1:])
	case "-pa":
		cmdPortFwdAdd(args[1:])
	case "-pl":
		requireArgs(args[1:], 0, "ipadm -pl")
		cmdPortFwdList()
	case "-pd":
		cmdPortFwdDelete(args[1:])
	default:
		if strings.HasPrefix(args[0], "-") {
			fail("unknown option %q\n\n%s", args[0], usage)
		}
		requireArgs(args[1:], 0, "ipadm <hostname>")
		cmdInteractive(args[0])
	}
}

func requireArgs(rest []string, want int, form string) {
	if len(rest) != want {
		fail("wrong number of arguments for %q\n\n%s", form, usage)
	}
}

// withStore opens the locked host DB, loads it, lets fn mutate it, and saves
// on success. fn returns the new host list plus a human-readable summary of
// what changed (printed on success).
func withStore(fn func(hosts []Host) ([]Host, string, error)) {
	s, err := openHostStore(dbPath())
	if err != nil {
		fail("%s", err)
	}
	defer s.Close()

	hosts, err := s.Load()
	if err != nil {
		fail("%s", err)
	}

	newHosts, summary, err := fn(hosts)
	if err != nil {
		fail("%s", err)
	}

	sortHosts(newHosts)
	if err := s.Save(newHosts); err != nil {
		fail("cannot write %s: %s", dbPath(), err)
	}
	fmt.Println(summary)
	fmt.Println("Note: run `ipadm -u` to apply this to dnsmasq/nftables.")
}

// withPortFwdStore mirrors withStore for the port-forward DB. It also loads
// the (read-only) host DB so add/delete can validate against known hosts.
func withPortFwdStore(fn func(hosts []Host, fwds []PortForward) ([]PortForward, string, error)) {
	hs, err := openHostStore(dbPath())
	if err != nil {
		fail("%s", err)
	}
	hosts, err := hs.Load()
	hs.Close()
	if err != nil {
		fail("%s", err)
	}

	s, err := openPortFwdStore(portFwdDBPath())
	if err != nil {
		fail("%s", err)
	}
	defer s.Close()

	fwds, err := s.Load()
	if err != nil {
		fail("%s", err)
	}

	newFwds, summary, err := fn(hosts, fwds)
	if err != nil {
		fail("%s", err)
	}

	sortPortFwds(newFwds)
	if err := s.Save(newFwds); err != nil {
		fail("cannot write %s: %s", portFwdDBPath(), err)
	}
	fmt.Println(summary)
	fmt.Println("Note: run `ipadm -u` to apply this to nftables.")
}

func cmdInteractive(name string) {
	cfg := loadNetConfig()
	withStore(func(hosts []Host) ([]Host, string, error) {
		newHosts, err := interactiveEdit(hosts, name, cfg)
		if err != nil {
			return nil, "", err
		}
		return newHosts, fmt.Sprintf("Host %q saved.", name), nil
	})
}

func cmdList() {
	s, err := openHostStore(dbPath())
	if err != nil {
		fail("%s", err)
	}
	defer s.Close()
	hosts, err := s.Load()
	if err != nil {
		fail("%s", err)
	}
	if len(hosts) == 0 {
		fmt.Println("no hosts configured")
		return
	}
	sortHosts(hosts)
	tw := tabwriter.NewWriter(os.Stdout, 2, 4, 2, ' ', 0)
	fmt.Fprintln(tw, "HOSTNAME\tIP\tMAC\tCOMMENT")
	for _, h := range hosts {
		mac := h.MAC
		if mac == "" {
			mac = "-"
		}
		fmt.Fprintf(tw, "%s\t%s\t%s\t%s\n", h.Name, h.IP, mac, h.Comment)
	}
	tw.Flush()
}

func cmdAdd(args []string) {
	force := false
	if len(args) > 0 && args[0] == "-f" {
		force = true
		args = args[1:]
	}
	if len(args) < 1 || len(args) > 3 {
		fail("usage: ipadm -a [-f] <hostname> [ip address] [mac address]")
	}
	name := args[0]
	ip, mac, err := classifyHostArgs(args[1:])
	if err != nil {
		fail("%s", err)
	}
	cfg := loadNetConfig()

	withStore(func(hosts []Host) ([]Host, string, error) {
		if err := validateHostname(name); err != nil {
			return nil, "", err
		}
		normMac, err := normalizeMAC(mac)
		if err != nil {
			return nil, "", err
		}

		idx := findHost(hosts, name)
		if idx >= 0 && !force {
			return nil, "", fmt.Errorf("host %q already exists (use -f to overwrite, or `ipadm %s` to edit it)", name, name)
		}

		cidr := cfg.cidrFor(normMac)
		if ip == "" {
			freeIP, err := nextFreeIP(hosts, cidr, idx, cfg.gatewayIP)
			if err != nil {
				return nil, "", err
			}
			ip = freeIP
		} else if err := validateIP(ip, cidr, cfg.gatewayIP); err != nil {
			return nil, "", err
		}

		if other := findByIP(hosts, ip, idx); other >= 0 {
			return nil, "", fmt.Errorf("IP %s is already assigned to host %q", ip, hosts[other].Name)
		}
		if other := findByMAC(hosts, normMac, idx); other >= 0 {
			return nil, "", fmt.Errorf("MAC %s is already assigned to host %q", normMac, hosts[other].Name)
		}

		h := Host{Name: name, IP: ip, MAC: normMac}
		if idx >= 0 {
			hosts[idx] = h
			return hosts, fmt.Sprintf("Host %q overwritten (-f), IP %s.", name, ip), nil
		}
		hosts = append(hosts, h)
		return hosts, fmt.Sprintf("Host %q added with IP %s.", name, ip), nil
	})
}

func cmdComment(args []string) {
	if len(args) < 2 {
		fail("usage: ipadm -c <hostname> <comment>")
	}
	name := args[0]
	comment := strings.Join(args[1:], " ")
	withStore(func(hosts []Host) ([]Host, string, error) {
		idx := findHost(hosts, name)
		if idx < 0 {
			return nil, "", fmt.Errorf("host %q not found", name)
		}
		hosts[idx].Comment = comment
		return hosts, fmt.Sprintf("Comment set for %q.", name), nil
	})
}

func cmdChangeIP(args []string) {
	if len(args) != 2 {
		fail("usage: ipadm -i <hostname> <ip address>")
	}
	name, ip := args[0], args[1]
	cfg := loadNetConfig()
	withStore(func(hosts []Host) ([]Host, string, error) {
		idx := findHost(hosts, name)
		if idx < 0 {
			return nil, "", fmt.Errorf("host %q not found", name)
		}
		cidr := cfg.cidrFor(hosts[idx].MAC)
		if err := validateIP(ip, cidr, cfg.gatewayIP); err != nil {
			return nil, "", err
		}
		if other := findByIP(hosts, ip, idx); other >= 0 {
			return nil, "", fmt.Errorf("IP %s is already assigned to host %q", ip, hosts[other].Name)
		}
		hosts[idx].IP = ip
		return hosts, fmt.Sprintf("IP of %q changed to %s.", name, ip), nil
	})
}

func cmdChangeMAC(args []string) {
	if len(args) != 2 {
		fail("usage: ipadm -m <hostname> <mac address>")
	}
	name, mac := args[0], args[1]
	if mac == "-" {
		mac = ""
	}
	cfg := loadNetConfig()
	withStore(func(hosts []Host) ([]Host, string, error) {
		idx := findHost(hosts, name)
		if idx < 0 {
			return nil, "", fmt.Errorf("host %q not found", name)
		}
		normMac, err := normalizeMAC(mac)
		if err != nil {
			return nil, "", err
		}
		if other := findByMAC(hosts, normMac, idx); other >= 0 {
			return nil, "", fmt.Errorf("MAC %s is already assigned to host %q", normMac, hosts[other].Name)
		}

		hosts[idx].MAC = normMac
		label := normMac
		if label == "" {
			label = "(removed)"
		}
		summary := fmt.Sprintf("MAC of %q changed to %s.", name, label)

		// The subnet a host belongs in depends on whether it has a MAC
		// (see netConfig). If that changed, move it to a free IP in the
		// now-correct subnet instead of leaving it in the wrong one.
		targetCIDR := cfg.cidrFor(normMac)
		if validateIP(hosts[idx].IP, targetCIDR, cfg.gatewayIP) != nil {
			newIP, err := nextFreeIP(hosts, targetCIDR, idx, cfg.gatewayIP)
			if err != nil {
				return nil, "", fmt.Errorf("MAC changed, but host must move to %s and no free IP is left there: %w", targetCIDR, err)
			}
			oldIP := hosts[idx].IP
			hosts[idx].IP = newIP
			summary += fmt.Sprintf(" Host moved from %s to %s (subnet %s).", oldIP, newIP, targetCIDR)
		}
		return hosts, summary, nil
	})
}

func cmdRename(args []string) {
	if len(args) != 2 {
		fail("usage: ipadm -r <old hostname> <new hostname>")
	}
	oldName, newName := args[0], args[1]
	withStore(func(hosts []Host) ([]Host, string, error) {
		if err := validateHostname(newName); err != nil {
			return nil, "", err
		}
		idx := findHost(hosts, oldName)
		if idx < 0 {
			return nil, "", fmt.Errorf("host %q not found", oldName)
		}
		if existing := findHost(hosts, newName); existing >= 0 && existing != idx {
			return nil, "", fmt.Errorf("host %q already exists", newName)
		}
		hosts[idx].Name = newName
		return hosts, fmt.Sprintf("Host %q renamed to %q.", oldName, newName), nil
	})
}

func cmdDelete(args []string) {
	if len(args) != 1 {
		fail("usage: ipadm -d <hostname>")
	}
	name := args[0]
	withStore(func(hosts []Host) ([]Host, string, error) {
		idx := findHost(hosts, name)
		if idx < 0 {
			return nil, "", fmt.Errorf("host %q not found", name)
		}
		hosts = append(hosts[:idx], hosts[idx+1:]...)
		return hosts, fmt.Sprintf("Host %q deleted.", name), nil
	})
}

func cmdPortFwdAdd(args []string) {
	force := false
	if len(args) > 0 && args[0] == "-f" {
		force = true
		args = args[1:]
	}
	if len(args) < 2 || len(args) > 4 {
		fail("usage: ipadm -pa [-f] <hostname> <wan-port> [lan-port] [tcp|udp|both]")
	}
	name := args[0]
	wanPort, err := parsePort(args[1])
	if err != nil {
		fail("%s", err)
	}

	withPortFwdStore(func(hosts []Host, fwds []PortForward) ([]PortForward, string, error) {
		if findHost(hosts, name) < 0 {
			return nil, "", fmt.Errorf("host %q is not in the ipadm database (add it first with `ipadm %s` or `ipadm -a`)", name, name)
		}
		lanPort, proto, err := classifyPortFwdArgs(args[2:], wanPort)
		if err != nil {
			return nil, "", err
		}
		if err := validatePort("wan-port", wanPort); err != nil {
			return nil, "", err
		}
		if err := validatePort("lan-port", lanPort); err != nil {
			return nil, "", err
		}

		conflicts := findConflictingPortFwds(fwds, wanPort, proto, -1)
		if len(conflicts) > 0 && !force {
			return nil, "", fmt.Errorf("wan-port %d/%s conflicts with an existing port-forward (use -f to overwrite)", wanPort, proto)
		}
		for i := len(conflicts) - 1; i >= 0; i-- {
			fwds = append(fwds[:conflicts[i]], fwds[conflicts[i]+1:]...)
		}

		fwds = append(fwds, PortForward{WanPort: wanPort, Proto: proto, Host: name, LanPort: lanPort})
		return fwds, fmt.Sprintf("Port-forward %d/%s -> %s:%d added.", wanPort, proto, name, lanPort), nil
	})
}

func cmdPortFwdList() {
	hs, err := openHostStore(dbPath())
	if err != nil {
		fail("%s", err)
	}
	hosts, err := hs.Load()
	hs.Close()
	if err != nil {
		fail("%s", err)
	}
	ipByHost := map[string]string{}
	for _, h := range hosts {
		ipByHost[strings.ToLower(h.Name)] = h.IP
	}

	s, err := openPortFwdStore(portFwdDBPath())
	if err != nil {
		fail("%s", err)
	}
	defer s.Close()
	fwds, err := s.Load()
	if err != nil {
		fail("%s", err)
	}
	if len(fwds) == 0 {
		fmt.Println("no port-forwards configured")
		return
	}
	sortPortFwds(fwds)
	tw := tabwriter.NewWriter(os.Stdout, 2, 4, 2, ' ', 0)
	fmt.Fprintln(tw, "WAN-PORT\tPROTO\tHOST\tLAN-IP\tLAN-PORT")
	for _, f := range fwds {
		ip, ok := ipByHost[strings.ToLower(f.Host)]
		if !ok {
			ip = "??? (host missing)"
		}
		fmt.Fprintf(tw, "%d\t%s\t%s\t%s\t%d\n", f.WanPort, f.Proto, f.Host, ip, f.LanPort)
	}
	tw.Flush()
}

func cmdPortFwdDelete(args []string) {
	if len(args) < 1 || len(args) > 2 {
		fail("usage: ipadm -pd <wan-port> [tcp|udp|both]")
	}
	wanPort, err := parsePort(args[0])
	if err != nil {
		fail("%s", err)
	}
	proto := ""
	if len(args) == 2 {
		proto = args[1]
		if !validProtos[proto] {
			fail("invalid protocol %q (expected tcp/udp/both)", proto)
		}
	}

	withPortFwdStore(func(hosts []Host, fwds []PortForward) ([]PortForward, string, error) {
		matches := findPortFwd(fwds, wanPort, proto)
		if len(matches) == 0 {
			return nil, "", fmt.Errorf("no port-forward found for wan-port %d", wanPort)
		}
		if len(matches) > 1 {
			var protos []string
			for _, i := range matches {
				protos = append(protos, fwds[i].Proto)
			}
			return nil, "", fmt.Errorf("wan-port %d is ambiguous (protocols: %s) — please specify a protocol", wanPort, strings.Join(protos, ", "))
		}
		idx := matches[0]
		deleted := fwds[idx]
		fwds = append(fwds[:idx], fwds[idx+1:]...)
		return fwds, fmt.Sprintf("Port-forward %d/%s (-> %s:%d) deleted.", deleted.WanPort, deleted.Proto, deleted.Host, deleted.LanPort), nil
	})
}

func parsePort(s string) (int, error) {
	var port int
	if _, err := fmt.Sscanf(s, "%d", &port); err != nil {
		return 0, fmt.Errorf("invalid port: %q", s)
	}
	return port, nil
}

func cmdUpdate() {
	cfg := loadNetConfig()

	hs, err := openHostStore(dbPath())
	if err != nil {
		fail("%s", err)
	}
	hosts, err := hs.Load()
	hs.Close()
	if err != nil {
		fail("%s", err)
	}

	dnsContent := renderDnsmasqConfig(hosts)
	if err := applyDnsmasqConfig(dnsmasqHostsPath(), dnsContent); err != nil {
		fail("%s", err)
	}
	fmt.Printf("%s updated, dnsmasq restarted (%d hosts).\n", dnsmasqHostsPath(), len(hosts))

	fs, err := openPortFwdStore(portFwdDBPath())
	if err != nil {
		fail("%s", err)
	}
	fwds, err := fs.Load()
	fs.Close()
	if err != nil {
		fail("%s", err)
	}

	nftContent, warnings := renderPortForwards(hosts, fwds, cfg.wanIface)
	for _, w := range warnings {
		fmt.Fprintln(os.Stderr, "warning:", w)
	}
	if err := applyPortForwards(nftPortFwdPath(), nftContent); err != nil {
		fail("%s", err)
	}
	fmt.Printf("%s updated, nftables reloaded (%d port-forwards).\n", nftPortFwdPath(), len(fwds)-len(warnings))
}