package wginterface

import (
	"fmt"
	"net"
	"net/netip"
	"os"
	"time"

	"golang.zx2c4.com/wireguard/wgctrl"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
)

const (
	// RekeyTimeout is the WireGuard session lifetime before a new handshake
	// is initiated. Sessions older than this but younger than SessionTimeout
	// remain valid.
	RekeyTimeout = 120 * time.Second

	// SessionTimeout is the WireGuard session lifetime after which sessions
	// are rejected. A peer with LastHandshakeTime older than this is
	// effectively disconnected.
	SessionTimeout = 180 * time.Second
)

const ProbeKeepalive = 8 * time.Second

var zeroKeepalive = time.Duration(0)

// Device wraps a wgctrl client bound to a named WireGuard interface.
type Device struct {
	client *wgctrl.Client
	name   string
}

// Open attaches to an existing WireGuard interface.
func Open(name string) (*Device, error) {
	client, err := wgctrl.New()
	if err != nil {
		return nil, fmt.Errorf("wgctrl: %w", err)
	}
	return &Device{client: client, name: name}, nil
}

// Close releases the underlying wgctrl client.
func (d *Device) Close() error {
	return d.client.Close()
}

// Name returns the interface name.
func (d *Device) Name() string {
	return d.name
}

// Configure sets the device's private key and UDP listen port.
func (d *Device) Configure(privKey wgtypes.Key, listenPort int) error {
	return d.client.ConfigureDevice(d.name, wgtypes.Config{
		PrivateKey: &privKey,
		ListenPort: &listenPort,
	})
}

// Peers returns the current state of all peers on the device.
func (d *Device) Peers() ([]wgtypes.Peer, error) {
	dev, err := d.client.Device(d.name)
	if err != nil {
		return nil, fmt.Errorf("get device %q: %w", d.name, err)
	}
	return dev.Peers, nil
}

// Peer returns the current state of a single peer by public key.
func (d *Device) Peer(pubKey wgtypes.Key) (wgtypes.Peer, error) {
	peers, err := d.Peers()
	if err != nil {
		return wgtypes.Peer{}, err
	}
	for _, p := range peers {
		if p.PublicKey == pubKey {
			return p, nil
		}
	}
	return wgtypes.Peer{}, fmt.Errorf("peer %v not found in %q", pubKey, d.name)
}

// AddPeer registers a peer with no AllowedIPs and no endpoint. WireGuard will
// accept handshakes from this peer but route no traffic to it yet.
func (d *Device) AddPeer(pubKey wgtypes.Key) error {
	return d.client.ConfigureDevice(d.name, wgtypes.Config{
		Peers: []wgtypes.PeerConfig{{
			PublicKey:         pubKey,
			ReplaceAllowedIPs: true,
		}},
	})
}

// SetRelay configures the relay peer with AllowedIPs covering the entire VPN
// network prefix. This is the fallback route for all VPN traffic.
func (d *Device) SetRelay(pubKey wgtypes.Key, endpoint netip.AddrPort, network netip.Prefix) error {
	masked := network.Masked()
	a4 := masked.Addr().As4()
	return d.client.ConfigureDevice(d.name, wgtypes.Config{
		Peers: []wgtypes.PeerConfig{{
			PublicKey: pubKey,
			Endpoint:  net.UDPAddrFromAddrPort(endpoint),
			AllowedIPs: []net.IPNet{{
				IP:   net.IP(a4[:]),
				Mask: net.CIDRMask(masked.Bits(), 32),
			}},
			ReplaceAllowedIPs: true,
		}},
	})
}

// AddProbe adds a peer with no AllowedIPs and an 8s keepalive. WireGuard will
// attempt handshakes without routing any traffic through this peer yet.
func (d *Device) AddProbe(pubKey wgtypes.Key, endpoint netip.AddrPort) error {
	keepalive := ProbeKeepalive
	return d.client.ConfigureDevice(d.name, wgtypes.Config{
		Peers: []wgtypes.PeerConfig{{
			PublicKey:                   pubKey,
			Endpoint:                    net.UDPAddrFromAddrPort(endpoint),
			AllowedIPs:                  []net.IPNet{},
			ReplaceAllowedIPs:           true,
			PersistentKeepaliveInterval: &keepalive,
		}},
	})
}

// Promote upgrades a probe entry to a /32 AllowedIPs and removes the probe
// keepalive, causing WireGuard to prefer this peer's direct path over the
// relay's wider route.
func (d *Device) Promote(pubKey wgtypes.Key, vpnIP netip.Addr) error {
	a4 := vpnIP.As4()
	return d.client.ConfigureDevice(d.name, wgtypes.Config{
		Peers: []wgtypes.PeerConfig{{
			PublicKey: pubKey,
			AllowedIPs: []net.IPNet{{
				IP:   net.IP(a4[:]),
				Mask: net.CIDRMask(32, 32),
			}},
			ReplaceAllowedIPs:           true,
			PersistentKeepaliveInterval: &zeroKeepalive,
		}},
	})
}

// AddDirect adds a peer with a known endpoint and /32 AllowedIPs in one step,
// for peers with a stable public endpoint reported by the hub.
func (d *Device) AddDirect(pubKey wgtypes.Key, endpoint netip.AddrPort, vpnIP netip.Addr) error {
	a4 := vpnIP.As4()
	return d.client.ConfigureDevice(d.name, wgtypes.Config{
		Peers: []wgtypes.PeerConfig{{
			PublicKey: pubKey,
			Endpoint:  net.UDPAddrFromAddrPort(endpoint),
			AllowedIPs: []net.IPNet{{
				IP:   net.IP(a4[:]),
				Mask: net.CIDRMask(32, 32),
			}},
			ReplaceAllowedIPs:           true,
			PersistentKeepaliveInterval: &zeroKeepalive,
		}},
	})
}

// RemovePeer removes a peer from the device.
func (d *Device) RemovePeer(pubKey wgtypes.Key) error {
	return d.client.ConfigureDevice(d.name, wgtypes.Config{
		Peers: []wgtypes.PeerConfig{{
			PublicKey: pubKey,
			Remove:    true,
		}},
	})
}

// EnableForwarding enables IPv4 forwarding globally and on the interface,
// required for relay peers that forward traffic between VPN peers.
func (d *Device) EnableForwarding() error {
	if err := os.WriteFile("/proc/sys/net/ipv4/ip_forward", []byte("1\n"), 0644); err != nil {
		return err
	}
	path := fmt.Sprintf("/proc/sys/net/ipv4/conf/%s/forwarding", d.name)
	return os.WriteFile(path, []byte("1\n"), 0644)
}