//go:build ignore

package peer

import (
	"log"
	"net/netip"
	"sync"
	"time"

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

const (
	probeWait    = 15 * time.Second
	backoffStart = 5 * time.Minute
	backoffMax   = time.Hour
)

type probeState struct {
	pubKey    wgtypes.Key
	endpoint  netip.AddrPort
	probing   bool
	direct    bool
	backoff   time.Duration
	backoffAt time.Time
}

type HolePunch struct {
	Globals
	mu    sync.Mutex
	peers [256]*probeState
}

func NewHolePunch(g Globals) *HolePunch {
	return &HolePunch{Globals: g}
}

// OnEndpointLearned is called when a peer's external WG endpoint becomes known,
// either from the VPN control channel (MsgMyEndpoint) or from the hub poller.
// fromHub=true resets any existing backoff so the hub-reported change is acted
// on immediately.
func (hp *HolePunch) OnEndpointLearned(peerIP byte, pubKey wgtypes.Key, endpoint netip.AddrPort, fromHub bool) {
	hp.mu.Lock()
	defer hp.mu.Unlock()

	ps := hp.peers[peerIP]
	if ps == nil {
		ps = &probeState{pubKey: pubKey}
		hp.peers[peerIP] = ps
	}

	if fromHub {
		ps.backoff = 0
		ps.backoffAt = time.Time{}
		ps.direct = false
		ps.pubKey = pubKey
	}
	ps.endpoint = endpoint

	if ps.probing || ps.direct {
		return
	}
	if ps.backoff > 0 && time.Now().Before(ps.backoffAt) {
		return
	}

	ps.probing = true
	go hp.runProbe(peerIP)
}

func (hp *HolePunch) runProbe(peerIP byte) {
	hp.mu.Lock()
	ps := hp.peers[peerIP]
	if ps == nil {
		hp.mu.Unlock()
		return
	}
	pubKey := ps.pubKey
	endpoint := ps.endpoint
	hp.mu.Unlock()

	vpnIP := netip.AddrFrom4([4]byte{
		hp.Network[0], hp.Network[1], hp.Network[2], peerIP,
	})

	probeStart := time.Now()
	if err := addProbeEntry(hp.WGClient, hp.WGDevName, pubKey, endpoint); err != nil {
		log.Printf("[HolePunch] addProbeEntry peer %d: %v", peerIP, err)
		hp.finishProbe(peerIP, false)
		return
	}

	time.Sleep(probeWait)

	_, handshakeTime, err := getPeerEndpoint(hp.WGClient, hp.WGDevName, pubKey)
	if err == nil && handshakeTime.After(probeStart) {
		if err := promoteToDirect(hp.WGClient, hp.WGDevName, pubKey, vpnIP); err != nil {
			log.Printf("[HolePunch] promoteToDirect peer %d: %v", peerIP, err)
		}
		hp.finishProbe(peerIP, true)
		return
	}

	// Probe failed — remove entry and schedule backoff retry.
	if err := removePeerEntry(hp.WGClient, hp.WGDevName, pubKey); err != nil {
		log.Printf("[HolePunch] removePeerEntry peer %d: %v", peerIP, err)
	}

	hp.mu.Lock()
	ps = hp.peers[peerIP]
	var delay time.Duration
	if ps != nil {
		if ps.backoff == 0 {
			ps.backoff = backoffStart
		} else {
			ps.backoff = min(ps.backoff*2, backoffMax)
		}
		ps.backoffAt = time.Now().Add(ps.backoff)
		ps.probing = false
		delay = ps.backoff
	}
	hp.mu.Unlock()

	if delay > 0 {
		go func() {
			time.Sleep(delay)
			hp.retryProbe(peerIP)
		}()
	}
}

func (hp *HolePunch) finishProbe(peerIP byte, success bool) {
	hp.mu.Lock()
	defer hp.mu.Unlock()
	ps := hp.peers[peerIP]
	if ps == nil {
		return
	}
	ps.probing = false
	if success {
		ps.direct = true
		ps.backoff = 0
	}
}

func (hp *HolePunch) retryProbe(peerIP byte) {
	hp.mu.Lock()
	ps := hp.peers[peerIP]
	if ps == nil || ps.probing || ps.direct {
		hp.mu.Unlock()
		return
	}
	ps.probing = true
	hp.mu.Unlock()
	hp.runProbe(peerIP)
}