package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/netip"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"regexp"
	"slices"
	"strings"
	"syscall"
	"time"
)

// The DHCP clients' names in DNS. Kea hands out the leases and Unbound
// answers, and nothing connects them, so this does: `router-webui leases-dns`
// reads Kea's lease file and publishes `<hostname>.<domain>` (and the reverse
// record) in the running Unbound through unbound-control.
//
// A client chooses its own host name, so it is not trusted: only the first
// label is used, only under the router's local domain, and a name that the
// configuration publishes (hosts, DHCP reservations, the router itself, a
// wildcard or a forward zone it falls under) is never taken over. Of two
// clients claiming one name the newer lease wins.

// staticNames is what the running configuration publishes itself
// (/etc/router/dns-static.json, written by the router module).
type staticNames struct {
	Domain       string   `json:"domain"`
	Names        []string `json:"names"`        // full names with records
	Overrides    []string `json:"overrides"`    // wildcard bases: the name and everything below
	ForwardZones []string `json:"forwardZones"` // answered by other resolvers
}

// DynRecord is a published client name.
type DynRecord struct {
	Name    string `json:"name"` // without the trailing dot
	Address string `json:"address"`
}

// DynSkip is a client that was not published, and why.
type DynSkip struct {
	Hostname string `json:"hostname"`
	Address  string `json:"address"`
	Reason   string `json:"reason"`
}

var labelRe = regexp.MustCompile(`^[a-z0-9]([a-z0-9-]{0,61}[a-z0-9])?$`)

func under(name, zone string) bool {
	return name == zone || strings.HasSuffix(name, "."+zone)
}

// planDynamic decides which leases are published.
func planDynamic(leases []Lease, st staticNames, now time.Time) (map[string]DynRecord, []DynSkip) {
	want := map[string]DynRecord{}
	var skipped []DynSkip
	if st.Domain == "" {
		return want, skipped
	}
	domain := strings.ToLower(strings.TrimSuffix(st.Domain, "."))
	static := map[string]bool{}
	for _, n := range st.Names {
		static[strings.ToLower(strings.TrimSuffix(n, "."))] = true
	}
	var active []Lease
	for _, l := range leases {
		if l.Hostname != "" && l.Expires.After(now) {
			active = append(active, l)
		}
	}
	// Newest lease first: it wins a name.
	slices.SortStableFunc(active, func(a, b Lease) int {
		if c := b.Start.Compare(a.Start); c != 0 {
			return c
		}
		return strings.Compare(a.Address, b.Address)
	})
	for _, l := range active {
		label, _, _ := strings.Cut(strings.ToLower(strings.TrimSpace(l.Hostname)), ".")
		skip := func(reason string) {
			skipped = append(skipped, DynSkip{Hostname: l.Hostname, Address: l.Address, Reason: reason})
		}
		if !labelRe.MatchString(label) {
			skip("not a usable host name")
			continue
		}
		if a, err := netip.ParseAddr(l.Address); err != nil || !a.Is4() {
			skip("not an IPv4 address")
			continue
		}
		name := label + "." + domain
		switch {
		case static[name]:
			skip("the configuration has a record with this name")
		case slices.ContainsFunc(st.Overrides, func(z string) bool { return under(name, strings.ToLower(strings.TrimSuffix(z, "."))) }):
			skip("below a wildcard record")
		case slices.ContainsFunc(st.ForwardZones, func(z string) bool { return under(name, strings.ToLower(strings.TrimSuffix(z, "."))) }):
			skip("below a forwarded zone")
		default:
			if other, ok := want[name]; ok {
				skip("a newer lease (" + other.Address + ") has this name")
				continue
			}
			want[name] = DynRecord{Name: name, Address: l.Address}
		}
	}
	return want, skipped
}

// reverseName is the in-addr.arpa name of an IPv4 address.
func reverseName(addr string) string {
	a, err := netip.ParseAddr(addr)
	if err != nil || !a.Is4() {
		return ""
	}
	b := a.As4()
	return fmt.Sprintf("%d.%d.%d.%d.in-addr.arpa", b[3], b[2], b[1], b[0])
}

// controller runs unbound-control commands.
type controller interface {
	run(args ...string) (string, error)
}

type unboundControl struct{ socket, config string }

func (u unboundControl) run(args ...string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	pre := []string{"-s", u.socket}
	if u.config != "" { // else unbound-control reads /etc/unbound/unbound.conf
		pre = append(pre, "-c", u.config)
	}
	out, err := exec.CommandContext(ctx, "unbound-control", append(pre, args...)...).CombinedOutput()
	if err != nil {
		err = fmt.Errorf("%w: %s", err, strings.TrimSpace(string(out)))
	}
	return strings.TrimSpace(string(out)), err
}

const dynTTL = "300"

func publish(ctl controller, r DynRecord) error {
	if _, err := ctl.run("local_data", r.Name+".", dynTTL, "IN", "A", r.Address); err != nil {
		return err
	}
	if rev := reverseName(r.Address); rev != "" {
		_, err := ctl.run("local_data", rev+".", dynTTL, "IN", "PTR", r.Name+".")
		return err
	}
	return nil
}

func unpublish(ctl controller, r DynRecord) error {
	var first error
	for _, n := range []string{r.Name, reverseName(r.Address)} {
		if n == "" {
			continue
		}
		if _, err := ctl.run("local_data_remove", n+"."); err != nil && first == nil {
			first = err
		}
	}
	return first
}

// syncDynamic brings Unbound from what was published to what is wanted. With
// fresh set (a new Unbound, or a restarted daemon) everything published
// before is removed first and the wanted records are all added.
func syncDynamic(ctl controller, published, want map[string]DynRecord, fresh bool) (map[string]DynRecord, error) {
	now := map[string]DynRecord{}
	var firstErr error
	note := func(err error) {
		if err != nil && firstErr == nil {
			firstErr = err
		}
	}
	for name, r := range published {
		if w, ok := want[name]; fresh || !ok || w.Address != r.Address {
			note(unpublish(ctl, r))
		} else {
			now[name] = r
		}
	}
	for name, w := range want {
		if _, ok := now[name]; ok {
			continue
		}
		if err := publish(ctl, w); err != nil {
			note(err)
			continue
		}
		now[name] = w
	}
	return now, firstErr
}

// dynStatus is what the daemon leaves for the web UI.
type dynStatus struct {
	Published []DynRecord `json:"published"`
	Skipped   []DynSkip   `json:"skipped"`
	Updated   time.Time   `json:"updated"`
	Error     string      `json:"error,omitempty"`
}

func readDynStatus(dir string) (dynStatus, bool) {
	var st dynStatus
	data, err := os.ReadFile(filepath.Join(dir, "status.json"))
	if err != nil || json.Unmarshal(data, &st) != nil {
		return st, false
	}
	return st, true
}

func writeDynStatus(dir string, st dynStatus) {
	data, _ := json.Marshal(st)
	if err := writeFileAtomic(filepath.Join(dir, "status.json"), data, 0o644); err != nil {
		log.Printf("status: %v", err)
	}
}

func recordsOf(m map[string]DynRecord) []DynRecord {
	out := make([]DynRecord, 0, len(m))
	for _, r := range m {
		out = append(out, r)
	}
	slices.SortFunc(out, func(a, b DynRecord) int { return strings.Compare(a.Name, b.Name) })
	return out
}

func loadStatic(path string) staticNames {
	var st staticNames
	if data, err := os.ReadFile(path); err == nil {
		json.Unmarshal(data, &st)
	}
	return st
}

// unboundPID tells a restarted Unbound (which lost what was published) from
// the same one.
func unboundPID(ctl controller) (string, error) {
	out, err := ctl.run("status")
	if err != nil {
		return "", err
	}
	return parseUnboundPID(out)
}

var pidRe = regexp.MustCompile(`\(pid ([0-9]+)\)`)

// parseUnboundPID reads the last line of `unbound-control status`:
// "unbound (pid 1234) is running...".
func parseUnboundPID(status string) (string, error) {
	if m := pidRe.FindStringSubmatch(status); m != nil {
		return m[1], nil
	}
	return "", fmt.Errorf("no pid in unbound-control status")
}

// leasesDNS is `router-webui leases-dns`.
func leasesDNS(args []string) {
	fs := flag.NewFlagSet("leases-dns", flag.ExitOnError)
	leases := fs.String("leases", "/var/lib/kea/dhcp4.leases", "Kea's DHCPv4 lease file")
	staticPath := fs.String("static", "/etc/router/dns-static.json", "the names the configuration publishes")
	socket := fs.String("socket", "/run/unbound/unbound.ctl", "Unbound's control socket")
	config := fs.String("unbound-config", "", "Unbound's config file (default: unbound-control's own)")
	run := fs.String("run", "/run/router-leases-dns", "directory for the status the web UI shows")
	interval := fs.Duration("interval", 5*time.Second, "how often the leases are read")
	once := fs.Bool("once", false, "do one round and exit")
	fs.Parse(args)

	ctl := unboundControl{socket: *socket, config: *config}
	if err := os.MkdirAll(*run, 0o755); err != nil {
		log.Fatal(err)
	}
	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM, os.Interrupt)
	defer stop()

	// What a previous run published is removed first: the daemon may have
	// been restarted while Unbound kept its data.
	published := map[string]DynRecord{}
	if st, ok := readDynStatus(*run); ok {
		for _, r := range st.Published {
			published[r.Name] = r
		}
	}
	pid, fresh := "", true
	var lastErr string
	for {
		st := loadStatic(*staticPath)
		want, skipped := planDynamic(readLeasesAt(*leases, time.Now()), st, time.Now())
		var err error
		if cur, perr := unboundPID(ctl); perr != nil {
			err = perr
		} else {
			if cur != pid {
				fresh = true // a new Unbound has none of it
			}
			var now map[string]DynRecord
			if now, err = syncDynamic(ctl, published, want, fresh); err == nil {
				fresh, pid = false, cur
			}
			published = now
		}
		status := dynStatus{Published: recordsOf(published), Skipped: skipped, Updated: time.Now()}
		if err != nil {
			status.Error = err.Error()
			if status.Error != lastErr {
				log.Printf("unbound-control: %v", err)
			}
		}
		lastErr = status.Error
		writeDynStatus(*run, status)
		if *once {
			if err != nil {
				os.Exit(1)
			}
			return
		}
		select {
		case <-ctx.Done():
			// Leave nothing behind: the names would outlive the daemon.
			syncDynamic(ctl, published, nil, false)
			writeDynStatus(*run, dynStatus{Updated: time.Now()})
			return
		case <-time.After(*interval):
		}
	}
}

// readLeasesAt reads the active leases of the lease file (without the LAN
// names the dashboard adds).
func readLeasesAt(path string, now time.Time) []Lease {
	return readLeases(path, nil, now)
}