package rcmgr

import (
	"encoding/binary"
	"net"
	"net/netip"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestItLimits(t *testing.T) {
	t.Run("IPv4", func(t *testing.T) {
		ip, err := netip.ParseAddr("1.2.3.4")
		require.NoError(t, err)
		cl := newConnLimiter()
		cl.connLimitPerSubnetV4[0].ConnCount = 1
		require.True(t, cl.addConn(ip))

		// should fail the second time
		require.False(t, cl.addConn(ip))

		otherIP, err := netip.ParseAddr("1.2.3.5")
		require.NoError(t, err)
		require.True(t, cl.addConn(otherIP))
	})
	t.Run("IPv6", func(t *testing.T) {
		ip, err := netip.ParseAddr("1:2:3:4::1")
		require.NoError(t, err)
		cl := newConnLimiter()
		original := cl.connLimitPerSubnetV6[0].ConnCount
		cl.connLimitPerSubnetV6[0].ConnCount = 1
		defer func() {
			cl.connLimitPerSubnetV6[0].ConnCount = original
		}()
		require.True(t, cl.addConn(ip))

		// should fail the second time
		require.False(t, cl.addConn(ip))
		otherIPSameSubnet := netip.MustParseAddr("1:2:3:4::2")
		require.False(t, cl.addConn(otherIPSameSubnet))

		otherIP := netip.MustParseAddr("2:2:3:4::2")
		require.True(t, cl.addConn(otherIP))
	})

	t.Run("IPv6 with multiple limits", func(t *testing.T) {
		cl := newConnLimiter()
		for i := 0; i < defaultMaxConcurrentConns; i++ {
			ip := net.ParseIP("ff:2:3:4::1")
			binary.BigEndian.PutUint16(ip[14:], uint16(i))
			ipAddr := netip.MustParseAddr(ip.String())
			require.True(t, cl.addConn(ipAddr))
		}

		// Next one should fail
		ip := net.ParseIP("ff:2:3:4::1")
		binary.BigEndian.PutUint16(ip[14:], uint16(defaultMaxConcurrentConns+1))
		require.False(t, cl.addConn(netip.MustParseAddr(ip.String())))

		// But on a different root subnet should work
		otherIP := netip.MustParseAddr("ffef:2:3::1")
		require.True(t, cl.addConn(otherIP))

		// But too many on the next subnet limit will fail too
		for i := 0; i < defaultMaxConcurrentConns*8; i++ {
			ip := net.ParseIP("ffef:2:3:4::1")
			binary.BigEndian.PutUint16(ip[5:7], uint16(i))
			ipAddr := netip.MustParseAddr(ip.String())
			require.True(t, cl.addConn(ipAddr))
		}

		ip = net.ParseIP("ffef:2:3:4::1")
		binary.BigEndian.PutUint16(ip[5:7], uint16(defaultMaxConcurrentConns*8+1))
		ipAddr := netip.MustParseAddr(ip.String())
		require.False(t, cl.addConn(ipAddr))
	})

	t.Run("IPv4 with localhost", func(t *testing.T) {
		cl := &connLimiter{
			networkPrefixLimitV4: DefaultNetworkPrefixLimitV4,
			connLimitPerSubnetV4: []ConnLimitPerSubnet{
				{PrefixLength: 0, ConnCount: 1}, // 1 connection for the whole IPv4 space
			},
		}

		ip := netip.MustParseAddr("1.2.3.4")
		require.True(t, cl.addConn(ip))

		ip = netip.MustParseAddr("4.3.2.1")
		// should fail the second time, we only allow 1 connection for the whole IPv4 space
		require.False(t, cl.addConn(ip))

		ip = netip.MustParseAddr("127.0.0.1")
		// Succeeds because we defined an explicit limit for the loopback subnet
		require.True(t, cl.addConn(ip))
	})
}

func genIP(data *[]byte) (netip.Addr, bool) {
	if len(*data) < 1 {
		return netip.Addr{}, false
	}

	genIP6 := (*data)[0]&0x01 == 1
	bytesRequired := 4
	if genIP6 {
		bytesRequired = 16
	}

	if len((*data)[1:]) < bytesRequired {
		return netip.Addr{}, false
	}

	*data = (*data)[1:]
	ip, ok := netip.AddrFromSlice((*data)[:bytesRequired])
	*data = (*data)[bytesRequired:]
	return ip, ok
}

func FuzzConnLimiter(f *testing.F) {
	// The goal is to try to enter a state where the count is incorrectly 0
	f.Fuzz(func(t *testing.T, data []byte) {
		ips := make([]netip.Addr, 0, len(data)/5)
		for {
			ip, ok := genIP(&data)
			if !ok {
				break
			}
			ips = append(ips, ip)
		}

		cl := newConnLimiter()
		addedConns := make([]netip.Addr, 0, len(ips))
		for _, ip := range ips {
			if cl.addConn(ip) {
				addedConns = append(addedConns, ip)
			}
		}

		addedCount := 0
		for _, ip := range cl.ip4connsPerLimit {
			for _, count := range ip {
				addedCount += count
			}
		}
		for _, ip := range cl.ip6connsPerLimit {
			for _, count := range ip {
				addedCount += count
			}
		}
		for _, count := range cl.connsPerNetworkPrefixV4 {
			addedCount += count
		}
		for _, count := range cl.connsPerNetworkPrefixV6 {
			addedCount += count
		}
		if addedCount == 0 && len(addedConns) > 0 {
			t.Fatalf("added count: %d", addedCount)
		}

		for _, ip := range addedConns {
			cl.rmConn(ip)
		}

		leftoverCount := 0
		for _, ip := range cl.ip4connsPerLimit {
			for _, count := range ip {
				leftoverCount += count
			}
		}
		for _, ip := range cl.ip6connsPerLimit {
			for _, count := range ip {
				leftoverCount += count
			}
		}
		for _, count := range cl.connsPerNetworkPrefixV4 {
			addedCount += count
		}
		for _, count := range cl.connsPerNetworkPrefixV6 {
			addedCount += count
		}
		if leftoverCount != 0 {
			t.Fatalf("leftover count: %d", leftoverCount)
		}
	})
}

func TestSortedNetworkPrefixLimits(t *testing.T) {
	npLimits := []NetworkPrefixLimit{
		{
			Network: netip.MustParsePrefix("1.2.0.0/16"),
		},
		{
			Network: netip.MustParsePrefix("1.2.3.0/28"),
		},
		{
			Network: netip.MustParsePrefix("1.2.3.4/32"),
		},
	}
	npLimits = sortNetworkPrefixes(npLimits)
	sorted := []NetworkPrefixLimit{
		{
			Network: netip.MustParsePrefix("1.2.3.4/32"),
		},
		{
			Network: netip.MustParsePrefix("1.2.3.0/28"),
		},
		{
			Network: netip.MustParsePrefix("1.2.0.0/16"),
		},
	}
	require.EqualValues(t, sorted, npLimits)
}