package main

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
)

// ─── service helpers ─────────────────────────────────────────────────────────────

func serviceEnable(os *OSInfo, iface string) {
	switch os.Init {
	case "systemd":
		runOrDie("systemctl", "enable", "awg-quick@"+iface)
	case "openrc":
		if err := run("rc-update", "add", "awg-quick."+iface, "default"); err != nil {
			warn("Could not register awg-quick.%s — may need manual setup", iface)
		}
	}
}

func serviceStart(os *OSInfo, iface string) {
	switch os.Init {
	case "systemd":
		runOrDie("systemctl", "start", "awg-quick@"+iface)
	case "openrc":
		runOrDie("awg-quick", "up", iface)
	}
}

func serviceStop(os *OSInfo, iface string) {
	switch os.Init {
	case "systemd":
		runOrDie("systemctl", "stop", "awg-quick@"+iface)
	case "openrc":
		runOrDie("awg-quick", "down", iface)
	}
}

func serviceRestart(os *OSInfo, iface string) {
	serviceStop(os, iface)
	serviceStart(os, iface)
}

func serviceStatus(os *OSInfo, iface string) {
	switch os.Init {
	case "systemd":
		// Best-effort; a non-zero status is not an error here.
		run("systemctl", "status", "awg-quick@"+iface, "--no-pager")
	case "openrc":
		run("rc-service", "awg-quick."+iface, "status")
	}
}

// ─── live peer management ────────────────────────────────────────────────────────

// awgIfaceUp reports whether the AmneziaWG interface is currently up.
func awgIfaceUp(iface string) bool {
	return silent("awg", "show", iface)
}

// awgPeerAdd hot-adds a peer to the running interface (no restart needed).
func awgPeerAdd(iface, pubkey, psk, ip string) {
	if !awgIfaceUp(iface) {
		warn("Interface %s is down — peer will be active on next start", iface)
		return
	}
	// awg reads the preshared key from a file; use a short-lived temp file.
	pskFile, err := os.CreateTemp("", "awg-psk-*")
	if err != nil {
		die("Failed to create temp psk file: %v", err)
	}
	defer os.Remove(pskFile.Name())
	if _, err := pskFile.WriteString(psk); err != nil {
		pskFile.Close()
		die("Failed to write temp psk file: %v", err)
	}
	pskFile.Close()

	if err := run("awg", "set", iface, "peer", pubkey,
		"preshared-key", pskFile.Name(),
		"allowed-ips", ip+"/32"); err != nil {
		// Non-fatal: by this point the client is already saved in the registry
		// and appended to <iface>.conf (createClient calls confAppendPeer first),
		// so the peer will pick up on the next restart/sync-config even if the
		// live hot-add fails. Dying here would report client creation as failed
		// when it actually succeeded, just without taking effect immediately.
		warn("awg set (peer add) failed: %v — peer saved, will apply on next restart/sync", err)
		return
	}
	info("Peer added live to %s (%s)", iface, ip)
}

// awgPeerRemove hot-removes a peer from the running interface.
func awgPeerRemove(iface, pubkey string) {
	if !awgIfaceUp(iface) {
		return
	}
	if err := run("awg", "set", iface, "peer", pubkey, "remove"); err != nil {
		die("awg set (peer remove) failed: %v", err)
	}
	info("Peer removed live from %s", iface)
}

// ─── config-file assembly ────────────────────────────────────────────────────────

func awgConfPath(iface string) string {
	return filepath.Join(awgConfDir, iface+".conf")
}

// awgConfHeader builds the [Interface] block including obfuscation parameters.
// PostUp loads the pre-generated nft ruleset; PostDown drops the table.
func awgConfHeader(c *Config, priv, srvIP, port, mtu string) string {
	return fmt.Sprintf(`[Interface]
PrivateKey = %s
Address = %s/24
ListenPort = %s
MTU = %s
SaveConfig = false

Jc = %s
Jmin = %s
Jmax = %s
S1 = %s
S2 = %s
H1 = %s
H2 = %s
H3 = %s
H4 = %s

PostUp = nft delete table inet awg_%%i 2>/dev/null || true
PostUp = nft -f %s/%%i-rules.nft
PostDown = nft delete table inet awg_%%i
`, priv, srvIP, port, mtu,
		c.Jc, c.Jmin, c.Jmax, c.S1, c.S2, c.H1, c.H2, c.H3, c.H4,
		awgConfDir)
}

// peerBlock renders a single [Peer] section for the interface config.
func peerBlock(name, pubkey, psk, ip string) string {
	return "\n[Peer]\n# " + name +
		"\nPublicKey = " + pubkey +
		"\nPresharedKey = " + psk +
		"\nAllowedIPs = " + ip + "/32\n"
}

// confAppendPeer appends a single [Peer] block to <conf-dir>/<iface>.conf.
func confAppendPeer(iface, name, pubkey, psk, ip string) {
	confPath := awgConfPath(iface)
	if _, err := os.Stat(confPath); err != nil {
		warn("%s not found — skipping conf update", confPath)
		return
	}
	f, err := os.OpenFile(confPath, os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		warn("%s not found — skipping conf update", confPath)
		return
	}
	defer f.Close()
	if _, err := f.WriteString(peerBlock(name, pubkey, psk, ip)); err != nil {
		die("Failed to append peer to %s: %v", confPath, err)
	}
	info("Peer appended to %s", confPath)
}

// confRebuild silently regenerates <conf-dir>/<iface>.conf from the registry
// (ACTIVE peers only) and rewrites the companion nft ruleset. Used by
// delete/enable/disable to keep the conf in sync without a restart prompt.
func confRebuild(c *Config, osInfo *OSInfo, clients []Client) {
	iface := c.Interface
	confPath := awgConfPath(iface)
	if _, err := os.Stat(confPath); err != nil {
		warn("%s not found — skipping conf rebuild", confPath)
		return
	}

	srvIP := serverIP(c.Network)
	defIface := defaultRouteIface()
	mtu := c.mtuOr("1420")

	var b strings.Builder
	b.WriteString(awgConfHeader(c, c.ServerPriv, srvIP, c.Port, mtu))
	for _, cl := range clients {
		if cl.IsEnabled == "ACTIVE" {
			b.WriteString(peerBlock(cl.Name, cl.PublicKey, cl.PSKKey, cl.IP))
		}
	}
	if err := os.WriteFile(confPath, []byte(b.String()), 0600); err != nil {
		die("Failed to write %s: %v", confPath, err)
	}
	writeNftRules(iface, defIface, mtu)
}

// serverIP derives the server's .1 host address from the network CIDR.
func serverIP(network string) string {
	base := networkBase(network)
	parts := strings.Split(base, ".")
	if len(parts) != 4 {
		die("Invalid network: %s", network)
	}
	return fmt.Sprintf("%s.%s.%s.1", parts[0], parts[1], parts[2])
}

// ─── nftables ruleset ────────────────────────────────────────────────────────────

// writeNftRules writes <conf-dir>/<iface>-rules.nft, loaded by PostUp via `nft -f`.
func writeNftRules(iface, defIface, mtu string) {
	mss := atoiOrDie(mtu) - 40
	nftPath := filepath.Join(awgConfDir, iface+"-rules.nft")

	content := fmt.Sprintf(`table inet awg_%s {
    # Allow forwarded traffic through the WG tunnel in both directions
    chain forward {
        type filter hook forward priority 0; policy accept;
        iif "%s" accept
        oif "%s" accept
    }

    # Masquerade only WG→external flows (more targeted than a blanket POSTROUTING rule)
    chain postrouting {
        type nat hook postrouting priority 100;
        iif "%s" oif "%s" masquerade
    }

    # MSS clamping keeps TCP segments within the WG tunnel MTU;
    # TTL normalisation hides the extra forwarding hop from remote hosts.
    chain mangle {
        type filter hook forward priority -150;
        iif "%s" tcp flags syn / syn,rst tcp option maxseg size set %d
        oif "%s" tcp flags syn / syn,rst tcp option maxseg size set %d
        iif "%s" oif "%s" ip ttl set 64
    }
}
`, iface, iface, iface, iface, defIface, iface, mss, iface, mss, iface, defIface)

	if err := os.WriteFile(nftPath, []byte(content), 0600); err != nil {
		die("Failed to write nft ruleset: %v", err)
	}
	info("NFT ruleset written: %s", nftPath)
}
