// Package wg wraps wireguard-go with a userspace network stack (gVisor netstack),
// the same approach Newt uses: no root, no kernel interface, works on Windows.
//
// A Device has one tunnel address. Connections made with DialContext and
// listeners made with ListenTCP/ListenUDP live entirely inside the tunnel.
package wg

import (
	"context"
	"crypto/ecdh"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"net"
	"net/netip"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.zx2c4.com/wireguard/conn"
	"golang.zx2c4.com/wireguard/device"
	"golang.zx2c4.com/wireguard/tun/netstack"
)

const (
	// MTU is conservative so the tunnel works over typical consumer links.
	MTU = 1280
	// Keepalive keeps NAT mappings open for peers behind NAT.
	Keepalive = 25
)

// GenerateKey returns a new base64-encoded private key.
func GenerateKey() (string, error) {
	k, err := ecdh.X25519().GenerateKey(rand.Reader)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(k.Bytes()), nil
}

// PublicKey derives the base64 public key from a base64 private key.
func PublicKey(priv string) (string, error) {
	raw, err := base64.StdEncoding.DecodeString(priv)
	if err != nil || len(raw) != 32 {
		return "", errors.New("invalid WireGuard private key")
	}
	k, err := ecdh.X25519().NewPrivateKey(raw)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(k.PublicKey().Bytes()), nil
}

func toHex(b64 string) (string, error) {
	raw, err := base64.StdEncoding.DecodeString(b64)
	if err != nil || len(raw) != 32 {
		return "", errors.New("invalid WireGuard key")
	}
	return hex.EncodeToString(raw), nil
}

// Peer is one remote end of the tunnel.
type Peer struct {
	PublicKey string // base64
	AllowedIP string // the peer's tunnel address, e.g. 100.89.0.7
	Endpoint  string // host:port, or "" when the peer connects to us
}

func (p Peer) sig() string { return p.PublicKey + "|" + p.AllowedIP + "|" + p.Endpoint }

type Device struct {
	dev *device.Device
	net *netstack.Net
	ip  netip.Addr

	mu    sync.Mutex
	peers map[string]Peer // by public key
}

// New brings up a tunnel endpoint at addr. listenPort 0 picks a random UDP port.
func New(privateKey string, addr netip.Addr, listenPort int) (*Device, error) {
	privHex, err := toHex(privateKey)
	if err != nil {
		return nil, err
	}
	tdev, tnet, err := netstack.CreateNetTUN([]netip.Addr{addr}, nil, MTU)
	if err != nil {
		return nil, err
	}
	dev := device.NewDevice(tdev, conn.NewDefaultBind(), device.NewLogger(device.LogLevelError, "wireguard: "))
	cfg := "private_key=" + privHex + "\n"
	if listenPort > 0 {
		cfg += "listen_port=" + strconv.Itoa(listenPort) + "\n"
	}
	if err := dev.IpcSet(cfg); err != nil {
		dev.Close()
		return nil, fmt.Errorf("configure wireguard: %w", err)
	}
	if err := dev.Up(); err != nil {
		dev.Close()
		return nil, err
	}
	return &Device{dev: dev, net: tnet, ip: addr, peers: map[string]Peer{}}, nil
}

func (d *Device) Addr() netip.Addr { return d.ip }

func (d *Device) Close() { d.dev.Close() }

// resolve turns host:port into the ip:port form wireguard-go requires.
func resolve(endpoint string) (string, error) {
	if endpoint == "" {
		return "", nil
	}
	a, err := net.ResolveUDPAddr("udp", endpoint)
	if err != nil {
		return "", err
	}
	return a.String(), nil
}

// SetPeers makes the device's peer list match want. Peers that are unchanged
// keep their sessions, so adding one site never interrupts the others.
// Peers whose endpoint can't be resolved are skipped (and returned in err).
func (d *Device) SetPeers(want []Peer) error {
	d.mu.Lock()
	defer d.mu.Unlock()
	var cfg strings.Builder
	var bad []string
	next := map[string]Peer{}
	for _, p := range want {
		ep, err := resolve(p.Endpoint)
		if err != nil {
			bad = append(bad, fmt.Sprintf("%s: %v", p.Endpoint, err))
			if old, ok := d.peers[p.PublicKey]; ok {
				next[p.PublicKey] = old // keep what works
			}
			continue
		}
		p.Endpoint = ep
		next[p.PublicKey] = p
	}
	for key := range d.peers {
		if _, keep := next[key]; !keep {
			h, _ := toHex(key)
			cfg.WriteString("public_key=" + h + "\nremove=true\n")
		}
	}
	keys := make([]string, 0, len(next))
	for k := range next {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, key := range keys {
		p := next[key]
		if old, ok := d.peers[key]; ok && old.sig() == p.sig() {
			continue
		}
		h, err := toHex(p.PublicKey)
		if err != nil {
			continue
		}
		cfg.WriteString("public_key=" + h + "\nreplace_allowed_ips=true\nallowed_ip=" + p.AllowedIP + "/32\n")
		if p.Endpoint != "" {
			cfg.WriteString("endpoint=" + p.Endpoint + "\npersistent_keepalive_interval=" + strconv.Itoa(Keepalive) + "\n")
		}
	}
	if cfg.Len() > 0 {
		if err := d.dev.IpcSet(cfg.String()); err != nil {
			return err
		}
	}
	d.peers = next
	if len(bad) > 0 {
		return errors.New(strings.Join(bad, "; "))
	}
	return nil
}

// Handshakes returns the time of the last completed handshake per peer public key.
func (d *Device) Handshakes() map[string]time.Time {
	out := map[string]time.Time{}
	state, err := d.dev.IpcGet()
	if err != nil {
		return out
	}
	var cur string
	for _, line := range strings.Split(state, "\n") {
		k, v, ok := strings.Cut(line, "=")
		if !ok {
			continue
		}
		switch k {
		case "public_key":
			raw, _ := hex.DecodeString(v)
			cur = base64.StdEncoding.EncodeToString(raw)
		case "last_handshake_time_sec":
			if n, err := strconv.ParseInt(v, 10, 64); err == nil && n > 0 && cur != "" {
				out[cur] = time.Unix(n, 0)
			}
		}
	}
	return out
}

func (d *Device) DialContext(ctx context.Context, network, addr string) (net.Conn, error) {
	return d.net.DialContext(ctx, network, addr)
}

func (d *Device) ListenTCP(port int) (net.Listener, error) {
	return d.net.ListenTCP(&net.TCPAddr{Port: port})
}

func (d *Device) ListenUDP(port int) (net.PacketConn, error) {
	return d.net.ListenUDP(&net.UDPAddr{Port: port})
}