package configurer

import (
	"net"
	"net/netip"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	wgconn "golang.zx2c4.com/wireguard/conn"
	wgdevice "golang.zx2c4.com/wireguard/device"
	"golang.zx2c4.com/wireguard/tun/tuntest"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"

	"github.com/netbirdio/netbird/client/iface/bind"
)

// newTestUSPConfigurer builds a configurer over a real wireguard-go device backed by an
// in-memory TUN. The device stays down, so no socket is opened and no privileges are needed.
func newTestUSPConfigurer(t *testing.T) *WGUSPConfigurer {
	t.Helper()

	tun := tuntest.NewChannelTUN()
	dev := wgdevice.NewDevice(tun.TUN(), wgconn.NewDefaultBind(), wgdevice.NewLogger(wgdevice.LogLevelSilent, ""))
	t.Cleanup(dev.Close)

	c := NewUSPConfigurerNoUAPI(dev, "wgtest0", bind.NewActivityRecorder())

	key, err := wgtypes.GeneratePrivateKey()
	require.NoError(t, err, "generate device private key")
	require.NoError(t, c.ConfigureInterface(key.String(), 0), "configure test device")

	return c
}

// seedPeers adds count peers, each with a /32 overlay address, and returns their public keys.
func seedPeers(t *testing.T, c *WGUSPConfigurer, count int) []string {
	t.Helper()

	keys := make([]string, 0, count)
	for i := 0; i < count; i++ {
		priv, err := wgtypes.GeneratePrivateKey()
		require.NoError(t, err, "generate peer private key")
		pub := priv.PublicKey().String()

		addr := netip.PrefixFrom(netip.AddrFrom4([4]byte{100, 64, byte(i >> 8), byte(i)}), 32)
		require.NoError(t, c.UpdatePeer(pub, []netip.Prefix{addr}, 25*time.Second, nil, nil), "add peer")
		keys = append(keys, pub)
	}
	return keys
}

func peerAllowedIPs(t *testing.T, c *WGUSPConfigurer, peerKey string) []string {
	t.Helper()

	stats, err := c.FullStats()
	require.NoError(t, err, "read device stats")

	for _, p := range stats.Peers {
		if p.PublicKey != peerKey {
			continue
		}
		got := make([]string, 0, len(p.AllowedIPs))
		for _, ipNet := range p.AllowedIPs {
			got = append(got, ipNet.String())
		}
		return got
	}
	t.Fatalf("peer %s not found on device", peerKey)
	return nil
}

// TestRemoveEndpointAddressPreservesRoutedAllowedIPs covers the prefixes the route manager
// attaches to a routing peer through AddAllowedIP. Those are not known to the peer.Conn that
// triggers the endpoint removal, so dropping them here would silently blackhole every route
// behind that peer on each relay or ICE disconnect.
func TestRemoveEndpointAddressPreservesRoutedAllowedIPs(t *testing.T) {
	c := newTestUSPConfigurer(t)
	peerKey := seedPeers(t, c, 3)[1]

	routed := []netip.Prefix{
		netip.MustParsePrefix("10.20.0.0/16"),
		netip.MustParsePrefix("192.168.7.0/24"),
	}
	for _, prefix := range routed {
		require.NoError(t, c.AddAllowedIP(peerKey, prefix), "add routed prefix")
	}

	before := peerAllowedIPs(t, c, peerKey)
	require.Len(t, before, 3, "peer should hold its overlay address plus both routed prefixes")

	require.NoError(t, c.RemoveEndpointAddress(peerKey), "remove endpoint address")

	assert.ElementsMatch(t, before, peerAllowedIPs(t, c, peerKey),
		"allowed IPs must survive the endpoint removal unchanged")
}

// TestRemoveEndpointAddressDoesNotScaleWithPeerCount is the regression guard for the actual
// defect: clearing one peer's endpoint used to dump and parse the whole device, so its cost
// grew with the size of the network map. On a routing peer with thousands of peers that dump
// runs on every relay and ICE transition, under the interface lock.
func TestRemoveEndpointAddressDoesNotScaleWithPeerCount(t *testing.T) {
	measure := func(peerCount int) float64 {
		c := newTestUSPConfigurer(t)
		peerKey := seedPeers(t, c, peerCount)[peerCount/2]

		return testing.AllocsPerRun(5, func() {
			require.NoError(t, c.RemoveEndpointAddress(peerKey), "remove endpoint address")
		})
	}

	small := measure(64)
	large := measure(1024)

	assert.Less(t, large, small*2,
		"clearing one endpoint allocated %.0f objects with 1024 peers against %.0f with 64: the cost still scales with the peer count",
		large, small)
}

// TestRemoveEndpointAddressFallsBackToDevice covers a peer the store never saw, which is what
// an out-of-band reconfiguration of the device leaves behind. The device stays the source of
// truth in that case, so the allowed IPs must still be preserved.
func TestRemoveEndpointAddressFallsBackToDevice(t *testing.T) {
	c := newTestUSPConfigurer(t)
	peerKey := seedPeers(t, c, 3)[1]
	require.NoError(t, c.AddAllowedIP(peerKey, netip.MustParsePrefix("10.20.0.0/16")), "add routed prefix")

	before := peerAllowedIPs(t, c, peerKey)
	c.allowedIPs.reset()

	require.NoError(t, c.RemoveEndpointAddress(peerKey), "remove endpoint address")

	assert.ElementsMatch(t, before, peerAllowedIPs(t, c, peerKey),
		"allowed IPs recovered from the device must be preserved")

	recovered, ok := c.allowedIPs.get(mustParseKey(t, peerKey))
	assert.True(t, ok, "the fallback must seed the store so the next call skips the device dump")
	assert.Len(t, recovered, 2, "seeded prefixes")
}

func TestRemoveAllowedIPKeepsTheOtherPrefixes(t *testing.T) {
	c := newTestUSPConfigurer(t)
	peerKey := seedPeers(t, c, 3)[0]
	routed := netip.MustParsePrefix("10.20.0.0/16")
	require.NoError(t, c.AddAllowedIP(peerKey, routed), "add routed prefix")
	require.NoError(t, c.AddAllowedIP(peerKey, netip.MustParsePrefix("192.168.7.0/24")), "add routed prefix")

	require.NoError(t, c.RemoveAllowedIP(peerKey, routed), "remove routed prefix")

	assert.ElementsMatch(t, []string{"100.64.0.0/32", "192.168.7.0/24"}, peerAllowedIPs(t, c, peerKey),
		"only the removed prefix should be gone")

	assert.ErrorIs(t, c.RemoveAllowedIP(peerKey, routed), ErrAllowedIPNotFound,
		"removing a prefix that is no longer configured must be reported")
}

// TestAddAllowedIPOnAbsentPeerDoesNotResurrectIt covers the lazy connection window documented
// in #6863: AddAllowedIP is update-only, a silent no-op when the peer is absent, so it must not
// leave the store claiming prefixes the device never took. RemoveEndpointAddress re-adds a peer
// without update-only, so a phantom entry would create a peer the device had dropped, and a
// created peer would steal those allowed IPs from whichever peer legitimately holds them.
func TestAddAllowedIPOnAbsentPeerDoesNotResurrectIt(t *testing.T) {
	c := newTestUSPConfigurer(t)
	seedPeers(t, c, 2)

	priv, err := wgtypes.GeneratePrivateKey()
	require.NoError(t, err, "generate peer private key")
	absent := priv.PublicKey().String()

	require.NoError(t, c.AddAllowedIP(absent, netip.MustParsePrefix("10.20.0.0/16")),
		"update-only add on an absent peer is a silent no-op")

	stats, err := c.FullStats()
	require.NoError(t, err, "read device stats")
	require.Len(t, stats.Peers, 2, "the absent peer must not have been created by AddAllowedIP")

	assert.ErrorIs(t, c.RemoveEndpointAddress(absent), ErrPeerNotFound,
		"clearing the endpoint of a peer the device does not have must fail")

	stats, err = c.FullStats()
	require.NoError(t, err, "read device stats")
	assert.Len(t, stats.Peers, 2, "no peer may be created while clearing an endpoint")
}

// TestRemoveEndpointAddressDoesNotStealAPrefixFromAnotherPeer covers WireGuard's rule that an
// allowed IP belongs to exactly one peer: configuring a prefix on a peer takes it away from
// whichever peer held it before. UpdatePeer relies on that rule rather than removing the prefix
// from the previous holder itself, so a prefix handed over between peers must not come back.
func TestRemoveEndpointAddressDoesNotStealAPrefixFromAnotherPeer(t *testing.T) {
	c := newTestUSPConfigurer(t)
	keys := seedPeers(t, c, 2)
	peerA, peerB := keys[0], keys[1]
	routed := netip.MustParsePrefix("10.20.0.0/16")

	require.NoError(t, c.AddAllowedIP(peerA, routed), "give the prefix to A")
	require.Contains(t, peerAllowedIPs(t, c, peerA), routed.String(), "A must hold the prefix")

	// The route moves to B. The device takes it away from A on its own.
	require.NoError(t, c.AddAllowedIP(peerB, routed), "hand the prefix over to B")
	require.Contains(t, peerAllowedIPs(t, c, peerB), routed.String(), "B must hold the prefix")
	require.NotContains(t, peerAllowedIPs(t, c, peerA), routed.String(), "the device must have taken it from A")

	require.NoError(t, c.RemoveEndpointAddress(peerA), "clear A's endpoint")

	assert.NotContains(t, peerAllowedIPs(t, c, peerA), routed.String(),
		"clearing A's endpoint must not take the prefix back from B")
	assert.Contains(t, peerAllowedIPs(t, c, peerB), routed.String(),
		"B must still hold the prefix")
}

// TestPresharedKeyCreatedPeerTakesPartInPrefixHandover covers a peer created by a preshared
// key write rather than by a peer update. Rosenpass applies a peer's first key without
// updateOnly, which creates the peer on the device, so a store that ignored that operation
// would treat the peer as unknown and would not account for a prefix later handed over to it.
func TestPresharedKeyCreatedPeerTakesPartInPrefixHandover(t *testing.T) {
	c := newTestUSPConfigurer(t)
	peerA := seedPeers(t, c, 1)[0]
	routed := netip.MustParsePrefix("10.20.0.0/16")
	require.NoError(t, c.AddAllowedIP(peerA, routed), "give the prefix to A")

	priv, err := wgtypes.GeneratePrivateKey()
	require.NoError(t, err, "generate peer private key")
	peerB := priv.PublicKey().String()

	psk, err := wgtypes.GenerateKey()
	require.NoError(t, err, "generate preshared key")
	require.NoError(t, c.SetPresharedKey(peerB, psk, false), "a first key creates the peer")

	require.NoError(t, c.AddAllowedIP(peerB, routed), "hand the prefix over to B")
	require.Contains(t, peerAllowedIPs(t, c, peerB), routed.String(), "B must hold the prefix")

	require.NoError(t, c.RemoveEndpointAddress(peerA), "clear A's endpoint")

	assert.NotContains(t, peerAllowedIPs(t, c, peerA), routed.String(),
		"clearing A's endpoint must not take the prefix back from B")
	assert.Contains(t, peerAllowedIPs(t, c, peerB), routed.String(), "B must still hold the prefix")
}

// TestUpdatePeerDoesNotWidenAMappedPrefixOnTheDevice is the end to end form of the
// conversion: a v4-mapped prefix must not reach the device as a zero length allowed IP,
// which would route every v4 address to that peer.
func TestUpdatePeerDoesNotWidenAMappedPrefixOnTheDevice(t *testing.T) {
	c := newTestUSPConfigurer(t)

	priv, err := wgtypes.GeneratePrivateKey()
	require.NoError(t, err, "generate peer private key")
	peerKey := priv.PublicKey().String()

	mapped := netip.MustParsePrefix("::ffff:10.1.2.3/112")
	require.NoError(t, c.UpdatePeer(peerKey, []netip.Prefix{mapped}, 25*time.Second, nil, nil), "add peer")

	onDevice := peerAllowedIPs(t, c, peerKey)
	assert.NotContains(t, onDevice, "0.0.0.0/0", "the device must not be given a catch-all allowed IP")
	assert.Equal(t, []string{"10.1.0.0/16"}, onDevice, "the device holds the normalized prefix")

	recorded, ok := c.allowedIPs.get(mustParseKey(t, peerKey))
	require.True(t, ok, "the peer must be recorded")
	require.Len(t, recorded, 1, "one prefix recorded")
	assert.Equal(t, onDevice[0], recorded[0].String(), "device and store must agree")
}

// TestUpdatePeerWithAnUnusableEndpointTouchesNothing pins the ordering: the endpoint is
// parsed before the device is configured, so a failure cannot leave the device holding a
// peer that the store never learned about, with the prefix handover skipped along with it.
func TestUpdatePeerWithAnUnusableEndpointTouchesNothing(t *testing.T) {
	c := newTestUSPConfigurer(t)
	seedPeers(t, c, 2)

	priv, err := wgtypes.GeneratePrivateKey()
	require.NoError(t, err, "generate peer private key")
	peerKey := priv.PublicKey().String()

	// A three byte address has no textual form netip can parse back.
	endpoint := &net.UDPAddr{IP: net.IP{1, 2, 3}, Port: 51820}
	require.Error(t, c.UpdatePeer(peerKey, []netip.Prefix{netip.MustParsePrefix("10.30.0.0/16")},
		25*time.Second, endpoint, nil), "an unusable endpoint must fail the update")

	stats, err := c.FullStats()
	require.NoError(t, err, "read device stats")
	assert.Len(t, stats.Peers, 2, "the peer must not have reached the device")

	_, ok := c.allowedIPs.get(mustParseKey(t, peerKey))
	assert.False(t, ok, "the peer must not have been recorded either")
}

// TestRemovePeerKeepsTheRecordWhenTheDeviceRefuses covers a removal that never reached the
// device. A single peer removal is one write, so a failure leaves the peer on the device
// exactly as it was, and the record still describes it; dropping it would only force the
// next caller to read the whole device back for an answer it already had.
func TestRemovePeerKeepsTheRecordWhenTheDeviceRefuses(t *testing.T) {
	c := newTestUSPConfigurer(t)
	peerKey := seedPeers(t, c, 1)[0]
	require.NoError(t, c.AddAllowedIP(peerKey, netip.MustParsePrefix("10.20.0.0/16")), "add routed prefix")

	before, ok := c.allowedIPs.get(mustParseKey(t, peerKey))
	require.True(t, ok, "the peer must be recorded before the removal")
	require.Len(t, before, 2, "overlay address plus routed prefix")

	// A closed device refuses every write, which is the shape of any failed removal.
	c.device.Close()

	require.Error(t, c.RemovePeer(peerKey), "the removal must report the failure")

	after, ok := c.allowedIPs.get(mustParseKey(t, peerKey))
	require.True(t, ok, "a peer still on the device must stay recorded")
	assert.Equal(t, before, after, "the record must describe the peer the device kept")
}

// mustParseKey turns the textual key the configurer API takes into the form the store
// keys on.
func mustParseKey(t *testing.T, key string) wgtypes.Key {
	t.Helper()

	parsed, err := wgtypes.ParseKey(key)
	require.NoError(t, err, "parse peer key")
	return parsed
}