package vpn

// Adding, renaming and removing the devices that dial in.
//
// Only OpenWrt's UCI is written, because that is where the server here is configured and
// because a half-supported second path is worse than an honest refusal. A wg-quick host is
// reported as unmanaged and the buttons do not appear.
//
// Two things happen for every change, in the order that fails safest:
//
//   - Adding: the configuration is written first and committed, then the peer is applied to
//     the running tunnel. If applying fails the configuration is rolled back, so the two never
//     disagree silently — a peer that works only after the next reboot is the worst outcome
//     of the three.
//   - Removing: the running tunnel first, because that is the moment access actually ends,
//     then the configuration. If the second half fails the caller is told plainly that the
//     device is off now but would return on a reboot.
//
// Nothing here restarts the network. `wg set` changes one peer on a live interface and
// touches nothing else, which matters rather a lot when the interface belongs to the router
// the house is behind.

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/netip"
	"os"
	"os/exec"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

// NameRe is what a device may be called: a short slug, because it also names the config
// section it is stored in.
var NameRe = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9 ._-]{0,30}$`)

// sectionRe guards the identifier that goes into a uci path.
var sectionRe = regexp.MustCompile(`^[a-zA-Z0-9_]{1,48}$`)

const sectionPrefix = "wgpeer_"

var (
	ErrUnmanaged = errors.New("this tunnel is not configured through UCI, so Ham Lab will not edit it here")
	ErrNoRoom    = errors.New("no free address is left in this tunnel's range")
)

// NewPeer is what the caller asks for.
type NewPeer struct {
	Iface string `json:"-"`
	Name  string `json:"name"`
	// Address is the tunnel address to give the device, with or without a prefix. Empty means
	// the next free one.
	Address string `json:"address"`
	// Endpoint is what the device will dial: the house's name or address, and the port.
	Endpoint string `json:"endpoint"`
	// DNS is what the device should use while connected. Empty leaves the line out.
	DNS string `json:"dns"`
	// Split routes only the home networks through the tunnel instead of everything.
	Split bool `json:"split"`
}

// Created is a new device, and the one chance to read its configuration.
type Created struct {
	Section   string `json:"section"`
	Name      string `json:"name"`
	PublicKey string `json:"public_key"`
	Address   string `json:"address"`
	// Config is the complete client profile. It holds the private key, which exists nowhere
	// else: it is generated here, handed over once and never stored.
	Config string `json:"config"`
}

// uciPeer is one peer section as UCI holds it.
type uciPeer struct {
	Section    string
	Iface      string // the wireguard_<iface> type it belongs to
	Desc       string
	PublicKey  string
	AllowedIPs []string
}

// uciPeers reads every peer section. The type of a peer section names its interface
// ("wireguard_wg0"), which is how a peer is tied to a tunnel.
func uciPeers(ctx context.Context) ([]uciPeer, error) {
	if _, err := exec.LookPath("uci"); err != nil {
		return nil, ErrUnmanaged
	}
	b, err := exec.CommandContext(ctx, "uci", "-q", "show", "network").Output()
	if err != nil {
		return nil, err
	}
	by := map[string]*uciPeer{}
	var order []string
	for _, line := range strings.Split(string(b), "\n") {
		lhs, rhs, ok := strings.Cut(line, "=")
		if !ok || !strings.HasPrefix(lhs, "network.") {
			continue
		}
		parts := strings.Split(strings.TrimPrefix(lhs, "network."), ".")
		v := strings.Trim(rhs, "'")
		section := parts[0]
		if len(parts) == 1 { // the section's own type
			if !strings.HasPrefix(v, "wireguard_") {
				continue
			}
			by[section] = &uciPeer{Section: section, Iface: strings.TrimPrefix(v, "wireguard_")}
			order = append(order, section)
			continue
		}
		p := by[section]
		if p == nil {
			continue
		}
		switch parts[1] {
		case "description":
			p.Desc = v
		case "public_key":
			p.PublicKey = v
		case "allowed_ips":
			p.AllowedIPs = strings.Fields(strings.ReplaceAll(v, ",", " "))
		}
	}
	out := make([]uciPeer, 0, len(order))
	for _, s := range order {
		out = append(out, *by[s])
	}
	return out, nil
}

// Managed reports whether this tunnel's peers can be edited here, and labels each peer with
// the section that holds it so the caller has a safe identifier to act on.
func Managed(ctx context.Context, t *Tunnel) bool {
	peers, err := uciPeers(ctx)
	if err != nil {
		return false
	}
	found := false
	byKey := map[string]string{}
	for _, p := range peers {
		if p.Iface == t.Name {
			byKey[p.PublicKey] = p.Section
			found = true
		}
	}
	for i := range t.Peers {
		t.Peers[i].Section = byKey[t.Peers[i].PublicKey]
	}
	// A tunnel with no peers at all is still manageable if UCI is the thing configuring it.
	if !found {
		for _, line := range uciInterfaces(ctx) {
			if line == t.Name {
				return true
			}
		}
	}
	return found
}

// uciInterfaces is every interface UCI configures as wireguard.
func uciInterfaces(ctx context.Context) []string {
	var out []string
	if _, err := exec.LookPath("uci"); err != nil {
		return out
	}
	b, err := exec.CommandContext(ctx, "uci", "-q", "show", "network").Output()
	if err != nil {
		return out
	}
	for _, line := range strings.Split(string(b), "\n") {
		lhs, rhs, ok := strings.Cut(line, "=")
		if !ok || strings.Trim(rhs, "'") != "wireguard" {
			continue
		}
		parts := strings.Split(strings.TrimPrefix(lhs, "network."), ".")
		if len(parts) == 2 && parts[1] == "proto" {
			out = append(out, parts[0])
		}
	}
	return out
}

// NextAddress is the lowest address in the tunnel's range that nothing is using. The
// interface's own address is taken, as is every address any peer is allowed.
func NextAddress(t Tunnel) (string, error) {
	var base netip.Prefix
	for _, a := range t.Addrs {
		p, err := netip.ParsePrefix(a)
		if err != nil || !p.Addr().Is4() {
			continue
		}
		base = p
		break
	}
	if !base.IsValid() {
		return "", errors.New("this tunnel has no IPv4 range to allocate from")
	}
	taken := map[netip.Addr]bool{base.Addr(): true}
	for _, p := range t.Peers {
		for _, a := range p.AllowedIPs {
			if ip, err := netip.ParsePrefix(strings.TrimSpace(a)); err == nil {
				taken[ip.Addr()] = true
			} else if ip, err := netip.ParseAddr(strings.TrimSpace(a)); err == nil {
				taken[ip] = true
			}
		}
	}
	ip := base.Masked().Addr().Next() // .0 is the network
	for base.Contains(ip) {
		next := ip.Next()
		if !base.Contains(next) { // the last address in the range is the broadcast
			break
		}
		if !taken[ip] {
			return ip.String(), nil
		}
		ip = next
	}
	return "", ErrNoRoom
}

// slug turns a device name into a section name, which may hold rather less.
func slug(name string) string {
	var b strings.Builder
	for _, r := range strings.ToLower(name) {
		switch {
		case r >= 'a' && r <= 'z', r >= '0' && r <= '9':
			b.WriteRune(r)
		case r == ' ', r == '-', r == '.', r == '_':
			b.WriteByte('_')
		}
	}
	s := strings.Trim(b.String(), "_")
	if s == "" {
		s = "peer"
	}
	if len(s) > 24 {
		s = s[:24]
	}
	return s
}

// genKeys makes a private key, its public half and a pre-shared key. The private key is
// returned to be handed over once and is never written anywhere by this daemon.
func genKeys(ctx context.Context) (priv, pub, psk string, err error) {
	run := func(name string, stdin string, args ...string) (string, error) {
		cmd := exec.CommandContext(ctx, name, args...)
		if stdin != "" {
			cmd.Stdin = strings.NewReader(stdin)
		}
		out, err := cmd.Output()
		return strings.TrimSpace(string(out)), err
	}
	if priv, err = run("wg", "", "genkey"); err != nil {
		return "", "", "", fmt.Errorf("generating a key: %w", err)
	}
	if pub, err = run("wg", priv+"\n", "pubkey"); err != nil {
		return "", "", "", fmt.Errorf("deriving the public key: %w", err)
	}
	if psk, err = run("wg", "", "genpsk"); err != nil {
		return "", "", "", fmt.Errorf("generating a pre-shared key: %w", err)
	}
	return priv, pub, psk, nil
}

// AddPeer creates a device: it writes the peer into the configuration, applies it to the
// running tunnel, and returns the client profile once.
func AddPeer(ctx context.Context, t Tunnel, n NewPeer) (*Created, error) {
	if !NameRe.MatchString(strings.TrimSpace(n.Name)) {
		return nil, errors.New("a device name is 1–31 letters, digits, spaces, dots, dashes or underscores")
	}
	name := strings.TrimSpace(n.Name)
	if strings.TrimSpace(n.Endpoint) == "" {
		return nil, errors.New("the device needs an endpoint to dial: the house's address or name, and the port")
	}
	endpoint, err := checkEndpoint(n.Endpoint)
	if err != nil {
		return nil, err
	}
	addr := strings.TrimSpace(n.Address)
	if addr == "" {
		if addr, err = NextAddress(t); err != nil {
			return nil, err
		}
	}
	host, err := hostOf(addr)
	if err != nil {
		return nil, err
	}
	for _, p := range t.Peers {
		for _, a := range p.AllowedIPs {
			if strings.HasPrefix(strings.TrimSpace(a), host+"/") {
				return nil, fmt.Errorf("%s is already used by %s", host, firstNonEmpty(p.Name, "another device"))
			}
		}
	}
	section := sectionPrefix + slug(name)
	existing, err := uciPeers(ctx)
	if err != nil {
		return nil, ErrUnmanaged
	}
	taken := map[string]bool{}
	for _, p := range existing {
		taken[p.Section] = true
	}
	for i := 2; taken[section]; i++ {
		section = fmt.Sprintf("%s%s%d", sectionPrefix, slug(name), i)
	}
	if !sectionRe.MatchString(section) {
		return nil, errors.New("that name cannot be turned into a configuration section")
	}

	priv, pub, psk, err := genKeys(ctx)
	if err != nil {
		return nil, err
	}
	allowed := host + "/32"

	// the configuration first, so a peer can never be live but unrecorded
	set := [][]string{
		{"set", "network." + section + "=wireguard_" + t.Name},
		{"set", "network." + section + ".description=" + name},
		{"set", "network." + section + ".public_key=" + pub},
		{"set", "network." + section + ".preshared_key=" + psk},
		{"set", "network." + section + ".allowed_ips=" + allowed},
	}
	for _, args := range set {
		if err := uci(ctx, args...); err != nil {
			_ = uci(ctx, "revert", "network")
			return nil, fmt.Errorf("writing the configuration: %w", err)
		}
	}
	if err := uci(ctx, "commit", "network"); err != nil {
		_ = uci(ctx, "revert", "network")
		return nil, fmt.Errorf("saving the configuration: %w", err)
	}

	// then the running tunnel; if this fails the configuration goes back as it was
	if err := wgSet(ctx, t.Name, pub, psk, allowed); err != nil {
		_ = uci(ctx, "delete", "network."+section)
		_ = uci(ctx, "commit", "network")
		return nil, fmt.Errorf("applying it to the running tunnel: %w", err)
	}

	return &Created{
		Section: section, Name: name, PublicKey: pub, Address: host,
		Config: profile(priv, psk, t, host, endpoint, n.DNS, n.Split),
	}, nil
}

// EditPeer changes a device's label or the address it is allowed. The section keeps the name
// it was created with; the description is what is shown.
func EditPeer(ctx context.Context, t Tunnel, section, name, address string) error {
	if !sectionRe.MatchString(section) {
		return errors.New("that is not a device this can edit")
	}
	peers, err := uciPeers(ctx)
	if err != nil {
		return ErrUnmanaged
	}
	var me *uciPeer
	for i := range peers {
		if peers[i].Section == section && peers[i].Iface == t.Name {
			me = &peers[i]
		}
	}
	if me == nil {
		return errors.New("no such device on this tunnel")
	}
	if name = strings.TrimSpace(name); name != "" {
		if !NameRe.MatchString(name) {
			return errors.New("a device name is 1–31 letters, digits, spaces, dots, dashes or underscores")
		}
		if err := uci(ctx, "set", "network."+section+".description="+name); err != nil {
			_ = uci(ctx, "revert", "network")
			return err
		}
	}
	allowed := ""
	if address = strings.TrimSpace(address); address != "" {
		host, err := hostOf(address)
		if err != nil {
			_ = uci(ctx, "revert", "network")
			return err
		}
		for _, p := range t.Peers {
			if p.Section == section {
				continue
			}
			for _, a := range p.AllowedIPs {
				if strings.HasPrefix(strings.TrimSpace(a), host+"/") {
					_ = uci(ctx, "revert", "network")
					return fmt.Errorf("%s is already used by %s", host, firstNonEmpty(p.Name, "another device"))
				}
			}
		}
		allowed = host + "/32"
		if err := uci(ctx, "set", "network."+section+".allowed_ips="+allowed); err != nil {
			_ = uci(ctx, "revert", "network")
			return err
		}
	}
	if err := uci(ctx, "commit", "network"); err != nil {
		_ = uci(ctx, "revert", "network")
		return err
	}
	if allowed != "" { // the running tunnel has to be told, or the device keeps the old address
		if err := wgSet(ctx, t.Name, me.PublicKey, "", allowed); err != nil {
			return fmt.Errorf("saved, but the running tunnel refused the new address: %w", err)
		}
	}
	return nil
}

// DeletePeer removes a device. The running tunnel goes first, because that is the moment the
// device actually loses its way in.
func DeletePeer(ctx context.Context, t Tunnel, section string) error {
	if !sectionRe.MatchString(section) {
		return errors.New("that is not a device this can remove")
	}
	peers, err := uciPeers(ctx)
	if err != nil {
		return ErrUnmanaged
	}
	var me *uciPeer
	for i := range peers {
		if peers[i].Section == section && peers[i].Iface == t.Name {
			me = &peers[i]
		}
	}
	if me == nil {
		return errors.New("no such device on this tunnel")
	}
	if err := wgRemove(ctx, t.Name, me.PublicKey); err != nil {
		return fmt.Errorf("removing it from the running tunnel: %w", err)
	}
	if err := uci(ctx, "delete", "network."+section); err != nil {
		return fmt.Errorf("the device is disconnected, but removing it from the configuration failed, so it would come back on a reboot: %w", err)
	}
	if err := uci(ctx, "commit", "network"); err != nil {
		return fmt.Errorf("the device is disconnected, but saving the configuration failed, so it would come back on a reboot: %w", err)
	}
	return nil
}

// ---- the small dangerous bits, each in one place ----

func uci(ctx context.Context, args ...string) error {
	out, err := exec.CommandContext(ctx, "uci", args...).CombinedOutput()
	if err != nil {
		return fmt.Errorf("uci %s: %s", strings.Join(args, " "), strings.TrimSpace(string(out)))
	}
	return nil
}

// wgSet applies a peer to the running interface. The pre-shared key goes in through a file
// because that is the only way wg takes one, so the file is made readable by nobody else and
// removed immediately.
func wgSet(ctx context.Context, iface, pub, psk, allowed string) error {
	args := []string{"set", iface, "peer", pub}
	if psk != "" {
		f, err := os.CreateTemp("", "hamlab-psk")
		if err != nil {
			return err
		}
		defer os.Remove(f.Name())
		if err := f.Chmod(0o600); err != nil {
			f.Close()
			return err
		}
		if _, err := f.WriteString(psk + "\n"); err != nil {
			f.Close()
			return err
		}
		f.Close()
		args = append(args, "preshared-key", f.Name())
	}
	args = append(args, "allowed-ips", allowed)
	out, err := exec.CommandContext(ctx, "wg", args...).CombinedOutput()
	if err != nil {
		return fmt.Errorf("wg set: %s", strings.TrimSpace(string(out)))
	}
	return nil
}

func wgRemove(ctx context.Context, iface, pub string) error {
	out, err := exec.CommandContext(ctx, "wg", "set", iface, "peer", pub, "remove").CombinedOutput()
	if err != nil {
		return fmt.Errorf("wg set: %s", strings.TrimSpace(string(out)))
	}
	return nil
}

// profile is the client's own configuration file, which is the only place its private key
// ever appears.
func profile(priv, psk string, t Tunnel, addr, endpoint, dns string, split bool) string {
	prefix := 24
	for _, a := range t.Addrs {
		if p, err := netip.ParsePrefix(a); err == nil && p.Addr().Is4() {
			prefix = p.Bits()
			break
		}
	}
	allowed := "0.0.0.0/0"
	if split {
		nets := []string{}
		for _, a := range t.Addrs {
			if p, err := netip.ParsePrefix(a); err == nil && p.Addr().Is4() {
				nets = append(nets, p.Masked().String())
			}
		}
		nets = append(nets, "10.0.0.0/8", "192.168.0.0/16", "172.16.0.0/12")
		allowed = strings.Join(dedupe(nets), ", ")
	}
	var b strings.Builder
	fmt.Fprintf(&b, "[Interface]\nPrivateKey = %s\nAddress = %s/%d\n", priv, addr, prefix)
	if dns = strings.TrimSpace(dns); dns != "" {
		fmt.Fprintf(&b, "DNS = %s\n", dns)
	}
	fmt.Fprintf(&b, "\n[Peer]\nPublicKey = %s\n", t.PublicKey)
	if psk != "" {
		fmt.Fprintf(&b, "PresharedKey = %s\n", psk)
	}
	fmt.Fprintf(&b, "Endpoint = %s\nAllowedIPs = %s\n", endpoint, allowed)
	return b.String()
}

func dedupe(in []string) []string {
	seen := map[string]bool{}
	out := in[:0:0]
	for _, s := range in {
		if !seen[s] {
			seen[s] = true
			out = append(out, s)
		}
	}
	sort.Strings(out)
	return out
}

// hostOf takes an address with or without a prefix and returns the bare address.
func hostOf(s string) (string, error) {
	s = strings.TrimSpace(s)
	if p, err := netip.ParsePrefix(s); err == nil {
		return p.Addr().String(), nil
	}
	a, err := netip.ParseAddr(s)
	if err != nil {
		return "", fmt.Errorf("%q is not an address", s)
	}
	return a.String(), nil
}

// checkEndpoint makes sure what the device will dial looks like a host and a port.
func checkEndpoint(s string) (string, error) {
	s = strings.TrimSpace(s)
	host, port, err := net.SplitHostPort(s)
	if err != nil || host == "" {
		return "", fmt.Errorf("%q should be a name or address and a port, like home.example.com:51820", s)
	}
	if n, err := strconv.Atoi(port); err != nil || n < 1 || n > 65535 {
		return "", fmt.Errorf("%q is not a port", port)
	}
	return s, nil
}

func firstNonEmpty(a, b string) string {
	if a != "" {
		return a
	}
	return b
}