package nebula

import (
	"context"
	"encoding/binary"
	"log/slog"
	"math/rand/v2"
	"net/netip"
	"sync"
	"sync/atomic"
	"time"

	"github.com/slackhq/nebula/config"
	"github.com/slackhq/nebula/header"
	"github.com/slackhq/nebula/overlay"
)

// PMTUD PoC: discover the path MTU per-tunnel via authenticated probes that ride
// the existing crypto session. We follow RFC 8899 PLPMTUD: a binary search
// between a known-good floor and a configured ceiling, with N consecutive probe
// losses at a size treated as "doesn't fit." Confirmed PMTU is pushed to the
// overlay device, which on Linux installs a per-host route with the discovered
// MTU. The kernel then surfaces EMSGSIZE / PTB to apps writing to the tun.
//
// Probe payload format (request):
//
//	[magic uint32 BE][probeID uint32 BE][padding 0x00...]
//
// Reply is a small ack with the same magic and probeID and no padding. We do not
// verify the reverse-path MTU; only the forward direction matters for the
// receiver's MTU on the inside.

const (
	pmtudMagic uint32 = 0x504D5544 // 'P' 'M' 'U' 'D'
	pmtudFloor        = 1280       // IPv6 minimum payload, also a safe internet MTU floor

	// pmtudConverged is the bytes-tolerance for stopping the search.
	pmtudConverged = 8

	// pmtudMaxLoss matches RFC 8899 MAX_PROBES (default 3).
	pmtudMaxLoss = 3

	// pmtudProbeInterval is the time between probe ticks during the search phase.
	// Once a peer converges the wheel stops ticking it; re-validation is driven
	// by connection_manager via MaybeProbeAsTest at its natural test cadence.
	pmtudProbeInterval = 500 * time.Millisecond

	// pmtudWheelMax is the wheel's maximum supported scheduling duration. We
	// only ever schedule at pmtudProbeInterval today, but the wheel needs a
	// max greater than its tick to allocate its slot ring sensibly.
	pmtudWheelMax = 5 * time.Second

	// pmtudOverheadPessimistic assumes IPv6 underlay + relay framing:
	//   IPv6(40) + UDP(8) + outer nebula(16) + outer AEAD tag(16)
	// + inner nebula(16) + inner AEAD tag(16) = 112 bytes.
	// TODO: track underlay address family and per-peer relay state on the HostInfo
	// so the manager can use the actual overhead for that tunnel and recover the
	// 32 bytes we pessimistically give up on direct IPv6 paths and the 52 bytes on
	// direct IPv4 paths.
	pmtudOverheadPessimistic = 112

	// pmtudUnsupportedAfter is the number of consecutive lost probes (across any
	// sizes) without ever receiving a reply that we treat as evidence the peer
	// does not understand the MTUDProbeRequest subtype (i.e. it's running an
	// older nebula). After this many failures with everReplied=false we mark the
	// peer pmtud-unsupported and stop scheduling probes. K is small enough that
	// it fires before the binary search would naturally converge to floor (which
	// would otherwise be ~30 wasted probes), but large enough to absorb a few
	// transient probe losses on a path that's just starting to settle.
	pmtudUnsupportedAfter = 5
)

// pmtudPeer tracks the binary-search state for one tunnel.
type pmtudPeer struct {
	mu       sync.Mutex
	addr     netip.Addr
	localIdx uint32

	// low is the largest outer IP packet size we have a confirmed ack for.
	// high is the smallest size we believe fails (the search ceiling to start).
	low, high int

	// inFlightSize is the outer IP packet size of the probe currently awaiting
	// an ack. 0 means no probe in flight.
	inFlightSize int
	// inFlightID matches the probeID echoed in the reply.
	inFlightID uint32
	// losses counts consecutive failures at inFlightSize.
	losses int

	// firstProbe is true until we have sent the first probe of a search. The
	// first probe targets the ceiling directly (RFC 8899 permits this Search
	// Algorithm choice); operators who set tun.max_mtu typically have a path
	// that supports it, so we converge in one probe in the common case.
	firstProbe bool
	// everReplied is true once we have ever received any MTUDProbeReply from
	// this peer. Combined with consecutiveFailures, this lets us detect peers
	// that don't understand the new subtype and stop probing them.
	everReplied bool
	// consecutiveFailures counts probes lost without an intervening reply.
	// Resets to 0 on any successful reply.
	consecutiveFailures int
	// unsupported is set true once we conclude the peer doesn't speak PMTUD.
	// The manager skips probes for unsupported peers.
	unsupported bool

	// converged means we have a confirmed PMTU and are in the slow re-validation phase.
	converged bool
	// applied is the inner MTU we last pushed to the overlay device (0 if never).
	applied int
}

func (p *pmtudPeer) overhead() int {
	// TODO: branch on actual underlay family + relay state for this peer.
	return pmtudOverheadPessimistic
}

func (p *pmtudPeer) midpoint() int {
	return (p.low + p.high) / 2
}

type pmtudManager struct {
	intf   *Interface
	device overlay.Device

	// peers is keyed by HostInfo.localIndexId.
	peers sync.Map // map[uint32]*pmtudPeer

	wheel *LockingTimerWheel[uint32]

	// floor is the always-safe inner MTU (= tun.mtu). Per-peer routes start here
	// on tunnel-up so unprobed traffic is always small enough to fit. Stored as
	// atomic int64 so reload can update it without coordinating with the readers
	// in tick/HandleReply/OnTunnelUp.
	floor atomic.Int64
	// ceiling is the search ceiling expressed as an outer IP packet size, derived
	// from tun.max_mtu (which is the kernel's device MTU on the tun) plus our
	// pessimistic overhead. PMTUD will not probe larger than this.
	ceiling atomic.Int64

	enabled atomic.Bool

	l *slog.Logger
}

func newPMTUDManagerFromConfig(l *slog.Logger, c *config.C, device overlay.Device) *pmtudManager {
	m := &pmtudManager{
		device: device,
		wheel:  NewLockingTimerWheel[uint32](pmtudProbeInterval, pmtudWheelMax),
		l:      l,
	}
	c.RegisterReloadCallback(func(c *config.C) { m.reload(c, false) })
	m.reload(c, true)
	return m
}

// reload applies tun.mtu / tun.max_mtu changes to the manager. On the initial
// call (during construction) it just snapshots state; on a live reload it also
// transitions in-flight peers to match the new bounds: clearing per-peer routes
// when newly disabled, seeding peers from the hostmap and flipping DF on
// outside sockets when newly enabled, and rebounding existing searches in
// place when only the ceiling moved.
func (m *pmtudManager) reload(c *config.C, initial bool) {
	if !initial && !c.HasChanged("tun.mtu") && !c.HasChanged("tun.max_mtu") {
		return
	}

	floor := c.GetInt("tun.mtu", overlay.DefaultMTU)
	maxMTU := c.GetInt("tun.max_mtu", 0)

	enable := maxMTU > floor && m.device.SupportsPerPeerMTU()
	var ceiling int
	if enable {
		ceiling = maxMTU + pmtudOverheadPessimistic
	}

	if initial {
		m.floor.Store(int64(floor))
		m.ceiling.Store(int64(ceiling))
		m.enabled.Store(enable)
		switch {
		case enable:
			m.l.Info("pmtud enabled", "floor", floor, "ceiling", ceiling, "tun.max_mtu", maxMTU)
		case maxMTU > floor:
			m.l.Warn("pmtud disabled: this platform does not yet support per-peer MTU routes",
				"tun.max_mtu", maxMTU)
		}
		return
	}

	wasEnabled := m.enabled.Load()
	m.floor.Store(int64(floor))
	m.ceiling.Store(int64(ceiling))
	m.enabled.Store(enable)

	switch {
	case wasEnabled && !enable:
		m.disableLive(floor, maxMTU)
	case !wasEnabled && enable:
		m.enableLive(floor, ceiling, maxMTU)
	case wasEnabled && enable:
		m.reboundLive(floor, ceiling, maxMTU)
	}
}

// disableLive clears per-peer routes and drops all peer state. We do not
// disable DF on the outside sockets; once on, it stays on for the life of the
// process. Operators flipping pmtud off live get correct routing behavior; if
// they want the historical no-DF behavior back they need to restart.
func (m *pmtudManager) disableLive(floor, maxMTU int) {
	m.peers.Range(func(k, v any) bool {
		p := v.(*pmtudPeer)
		p.mu.Lock()
		applied := p.applied
		addr := p.addr
		p.applied = 0
		p.mu.Unlock()
		if applied != 0 {
			if err := m.device.SetPeerMTU(addr, 0); err != nil {
				m.l.Warn("pmtud: failed to clear per-peer mtu on disable", "addr", addr, "error", err)
			}
		}
		m.peers.Delete(k)
		return true
	})
	m.l.Info("pmtud disabled (tun.max_mtu <= tun.mtu)", "tun.mtu", floor, "tun.max_mtu", maxMTU)
}

// enableLive flips DF on every outside socket. We don't pre-seed existing
// tunnels here; connection_manager's normal test cadence will eventually call
// MaybeProbeAsTest for each peer, which seeds on miss and lets the wheel pick
// up the search from there. New tunnels established after this point still
// take the OnTunnelUp fast path.
func (m *pmtudManager) enableLive(floor, ceiling, maxMTU int) {
	m.enableDF()
	m.l.Info("pmtud enabled", "floor", floor, "ceiling", ceiling, "tun.max_mtu", maxMTU)
}

// reboundLive resets each peer's search state to the new bounds. Peers whose
// confirmed PMTU still fits under the new ceiling keep their applied route in
// place during the new search; peers whose confirmed PMTU exceeds the new
// ceiling get cleared back to floor and re-search from scratch. The unsupported
// flag is preserved because peer software version doesn't change on reload.
func (m *pmtudManager) reboundLive(floor, ceiling, maxMTU int) {
	overhead := pmtudOverheadPessimistic
	m.peers.Range(func(k, v any) bool {
		p := v.(*pmtudPeer)
		p.mu.Lock()
		if p.applied > 0 && p.applied+overhead > ceiling {
			if err := m.device.SetPeerMTU(p.addr, 0); err != nil {
				m.l.Warn("pmtud: failed to clear per-peer mtu on rebound", "addr", p.addr, "error", err)
			} else {
				p.applied = 0
			}
		}
		p.low = floor + overhead
		p.high = ceiling
		p.inFlightSize = 0
		p.inFlightID = 0
		p.losses = 0
		p.firstProbe = !p.unsupported
		p.converged = false
		idx := p.localIdx
		unsupported := p.unsupported
		p.mu.Unlock()
		if !unsupported {
			m.wheel.Add(idx, pmtudProbeInterval)
		}
		return true
	})
	m.l.Info("pmtud reloaded", "floor", floor, "ceiling", ceiling, "tun.max_mtu", maxMTU)
}

// enableDF asks every outside socket to set the don't-fragment bit on outbound
// packets. Idempotent: safe to call from both Start (initial enable) and from a
// live reload that flips pmtud on.
func (m *pmtudManager) enableDF() {
	for i, w := range m.intf.writers {
		if err := w.EnablePathMTUDiscovery(); err != nil {
			m.l.Warn("pmtud: failed to enable path mtu discovery on outside socket; pmtud will not work correctly",
				"writer", i, "error", err)
		}
	}
}

// Start runs the probe scheduler until ctx is done. The loop runs even when PMTUD
// is disabled at startup so a hot reload can turn it on without restarting nebula.
//
// When PMTUD is enabled at startup we ask each outside socket to enable
// path-MTU discovery (DF on every send). This is intentionally gated on the
// feature being on so that operators who haven't opted in keep the historical
// behavior where the kernel may fragment outbound nebula UDP packets. A live
// reload from disabled to enabled will also flip DF on via enableLive; the
// reverse direction does not turn DF off, so flipping pmtud back off live
// keeps DF on until restart.
func (m *pmtudManager) Start(ctx context.Context) {
	if m.enabled.Load() {
		m.enableDF()
	}

	ticker := time.NewTicker(m.wheel.t.tickDuration)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case now := <-ticker.C:
			m.wheel.Advance(now)
			for {
				idx, has := m.wheel.Purge()
				if !has {
					break
				}
				m.tick(idx)
			}
		}
	}
}

// OnTunnelUp is called when a HostInfo becomes traffic-watched. The kernel
// already routes packets to this peer through the per-vpn-network route (mtu =
// tun.mtu), so the floor is in effect implicitly. We just kick off the search
// here; HandleReply will install a per-host /32 (or /128) route once a larger
// size is confirmed.
func (m *pmtudManager) OnTunnelUp(hi *HostInfo) {
	if !m.enabled.Load() {
		return
	}
	m.seedPeer(hi)
}

// seedPeer is the shared body of OnTunnelUp and the live-reload enable path.
// LoadOrStore protects against double-seeding the same localIndexId from a
// race between OnTunnelUp and a reload-driven hostmap walk.
func (m *pmtudManager) seedPeer(hi *HostInfo) {
	if hi == nil || len(hi.vpnAddrs) == 0 {
		return
	}
	floor := int(m.floor.Load())
	ceiling := int(m.ceiling.Load())
	p := &pmtudPeer{
		addr:       hi.vpnAddrs[0],
		localIdx:   hi.localIndexId,
		low:        floor + pmtudOverheadPessimistic,
		high:       ceiling,
		firstProbe: true,
	}
	if _, loaded := m.peers.LoadOrStore(hi.localIndexId, p); loaded {
		return
	}
	m.wheel.Add(hi.localIndexId, pmtudProbeInterval)
}

// OnTunnelDown is called when a HostInfo is being torn down. Removes any per-host
// MTU override so the device default applies again.
func (m *pmtudManager) OnTunnelDown(hi *HostInfo) {
	if hi == nil {
		return
	}
	v, ok := m.peers.LoadAndDelete(hi.localIndexId)
	if !ok {
		return
	}
	p := v.(*pmtudPeer)
	p.mu.Lock()
	applied := p.applied
	addr := p.addr
	p.applied = 0
	p.mu.Unlock()
	if applied != 0 {
		if err := m.device.SetPeerMTU(addr, 0); err != nil {
			m.l.Warn("pmtud: failed to clear per-peer mtu", "addr", addr, "error", err)
		}
	}
}

// OnRoam is called when a HostInfo's remote underlay address changes. The path
// MTU may now be different; drop the per-host route so the kernel falls back to
// the per-vpn-network route (mtu = tun.mtu floor), then restart the search.
// We do not reset the unsupported flag: peer software version doesn't change on
// roam, so once we've decided a peer doesn't speak PMTUD we stay decided.
func (m *pmtudManager) OnRoam(hi *HostInfo) {
	if !m.enabled.Load() || hi == nil {
		return
	}
	v, ok := m.peers.Load(hi.localIndexId)
	if !ok {
		return
	}
	p := v.(*pmtudPeer)
	p.mu.Lock()
	if p.unsupported {
		p.mu.Unlock()
		return
	}
	p.low = int(m.floor.Load()) + pmtudOverheadPessimistic
	p.high = int(m.ceiling.Load())
	p.inFlightSize = 0
	p.inFlightID = 0
	p.losses = 0
	p.consecutiveFailures = 0
	p.firstProbe = true
	p.converged = false
	if p.applied != 0 {
		if err := m.device.SetPeerMTU(p.addr, 0); err != nil {
			m.l.Warn("pmtud: failed to clear per-peer mtu on roam", "addr", p.addr, "error", err)
		} else {
			p.applied = 0
		}
	}
	p.mu.Unlock()
	m.wheel.Add(hi.localIndexId, pmtudProbeInterval)
}

// MaybeProbeAsTest is called by connection_manager when it would otherwise send
// a TestRequest because a tunnel has gone silent. If we have a confirmed PMTU
// for this peer that's larger than the floor, we send a probe at that size
// instead. The reply confirms both liveness (consumed by connection_manager via
// the existing inbound traffic accounting fallthrough in outside.go) and that
// the confirmed PMTU still fits (consumed by HandleReply here). One synthetic
// packet does the work of two.
//
// Returns true if a probe was sent. False means the caller should send a
// regular TestRequest at the floor.
//
// On probe failure, connection_manager's existing pendingDeletion timeout will
// tear the tunnel down. Heavy hammer, but correct: a re-handshake re-runs PMTUD
// discovery against the now-shrunken path. A future EMSGSIZE-capture followup
// can replace this with a soft-drop-and-research flow.
func (m *pmtudManager) MaybeProbeAsTest(hi *HostInfo) bool {
	if !m.enabled.Load() || hi == nil {
		return false
	}
	v, ok := m.peers.Load(hi.localIndexId)
	if !ok {
		// Tunnel pre-dates the manager being aware of it (e.g. pmtud was just
		// enabled live, or AddTrafficWatch fired before this call). Seed the
		// peer so the wheel picks up the search; let connection_manager send
		// its regular TestRequest this cycle.
		m.seedPeer(hi)
		return false
	}
	p := v.(*pmtudPeer)
	p.mu.Lock()
	if p.unsupported || p.applied == 0 {
		p.mu.Unlock()
		return false
	}
	overhead := p.overhead()
	size := p.applied + overhead
	id := rand.Uint32()
	p.inFlightSize = size
	p.inFlightID = id
	p.mu.Unlock()

	m.sendProbe(hi, size, id, overhead)
	return true
}

// HandleReply consumes an MTUDProbeReply payload from the receive path.
func (m *pmtudManager) HandleReply(localIdx uint32, payload []byte) {
	if !m.enabled.Load() {
		return
	}
	if len(payload) < 8 {
		return
	}
	if binary.BigEndian.Uint32(payload[0:4]) != pmtudMagic {
		return
	}
	id := binary.BigEndian.Uint32(payload[4:8])

	v, ok := m.peers.Load(localIdx)
	if !ok {
		return
	}
	p := v.(*pmtudPeer)
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.inFlightSize == 0 || p.inFlightID != id {
		return
	}

	confirmed := p.inFlightSize
	p.low = confirmed
	p.inFlightSize = 0
	p.losses = 0
	p.everReplied = true
	p.consecutiveFailures = 0

	innerMTU := confirmed - p.overhead()
	// Only install a /32 override when it would actually raise the MTU above the
	// per-vpn-network floor route. If the discovered MTU is <= floor, the /24
	// already covers it; installing a /32 at floor would just create roam churn.
	if innerMTU > int(m.floor.Load()) && p.applied != innerMTU {
		if err := m.device.SetPeerMTU(p.addr, innerMTU); err != nil {
			m.l.Warn("pmtud: failed to apply per-peer mtu", "addr", p.addr, "innerMTU", innerMTU, "error", err)
		} else {
			m.l.Info("pmtud probe confirmed",
				"addr", p.addr,
				"outerMTU", confirmed,
				"innerMTU", innerMTU,
				"low", p.low,
				"high", p.high,
			)
			p.applied = innerMTU
		}
	}

	if p.high-p.low <= pmtudConverged {
		p.converged = true
	} else {
		p.converged = false
	}
}

// tick handles one wheel firing for a single peer.
func (m *pmtudManager) tick(localIdx uint32) {
	v, ok := m.peers.Load(localIdx)
	if !ok {
		return
	}
	p := v.(*pmtudPeer)
	p.mu.Lock()

	if p.unsupported {
		p.mu.Unlock()
		return
	}

	// If a probe was outstanding, this tick is the loss timeout.
	if p.inFlightSize != 0 {
		p.losses++
		p.consecutiveFailures++
		if p.losses >= pmtudMaxLoss {
			p.high = p.inFlightSize
			p.inFlightSize = 0
			p.losses = 0
			if p.high-p.low <= pmtudConverged {
				p.converged = true
			}
		}
	}

	// If we've never gotten a reply from this peer and we've burned through our
	// failure budget, conclude the peer doesn't understand the MTUDProbeRequest
	// subtype and stop scheduling probes for it.
	if !p.everReplied && p.consecutiveFailures >= pmtudUnsupportedAfter {
		p.unsupported = true
		addr := p.addr
		p.mu.Unlock()
		m.l.Info("pmtud: peer not responding to probes, marking unsupported",
			"addr", addr, "failures", pmtudUnsupportedAfter)
		return
	}

	hi := m.intf.hostMap.QueryIndex(localIdx)
	if hi == nil {
		p.mu.Unlock()
		m.peers.Delete(localIdx)
		return
	}

	// Once a peer converges, the wheel stops scheduling for it. Re-validation
	// (and the resulting black hole detection) is driven by connection_manager
	// via MaybeProbeAsTest at its natural test cadence, so a converged peer
	// has nothing for the wheel to do until OnRoam or a tunnel down/up cycle
	// triggers a fresh search.
	if p.converged {
		p.mu.Unlock()
		return
	}

	ceiling := int(m.ceiling.Load())
	var size int
	switch {
	case p.firstProbe:
		// Probe the ceiling directly. If the path supports it (the common case
		// when an operator has explicitly configured tun.max_mtu), we converge
		// in one round trip. If it fails, the standard binary search resumes
		// on the next tick from the (low, ceiling) bounds.
		size = ceiling
		p.firstProbe = false
	case p.losses > 0 && p.inFlightSize != 0:
		size = p.inFlightSize
	default:
		size = p.midpoint()
	}
	if size < pmtudFloor {
		size = pmtudFloor
	}
	if size > ceiling {
		size = ceiling
	}

	id := rand.Uint32()
	p.inFlightSize = size
	p.inFlightID = id
	overhead := p.overhead()
	p.mu.Unlock()

	m.sendProbe(hi, size, id, overhead)
	m.wheel.Add(localIdx, pmtudProbeInterval)
}

// sendProbe builds an MTUDProbeRequest payload that will produce an outer IP
// packet of approximately `outerSize` bytes, then sends it.
func (m *pmtudManager) sendProbe(hi *HostInfo, outerSize int, id uint32, overhead int) {
	payloadLen := outerSize - overhead
	if payloadLen < 8 {
		payloadLen = 8
	}
	p := make([]byte, payloadLen)
	binary.BigEndian.PutUint32(p[0:4], pmtudMagic)
	binary.BigEndian.PutUint32(p[4:8], id)
	// remaining bytes are zero-padding

	nb := make([]byte, 12)
	out := make([]byte, outerSize+128) // headroom for header/tag/relay framing
	m.intf.SendMessageToHostInfo(header.Test, header.MTUDProbeRequest, hi, p, nb, out)
}