// Copyright 2025 Antonio Ojea
// SPDX-License-Identifier: Apache-2.0

package network

import (
	"math/rand"
	"net/netip"
	"reflect"
	"sort"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"k8s.io/apimachinery/pkg/util/sets"
)

func Test_InsertGetDelete(t *testing.T) {
	testCases := []struct {
		name   string
		prefix netip.Prefix
	}{
		{
			name:   "ipv4",
			prefix: netip.MustParsePrefix("192.168.0.0/24"),
		},
		{
			name:   "ipv6",
			prefix: netip.MustParsePrefix("fd00:1:2:3::/124"),
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			tree := NewIPTree[int]()
			ok := tree.InsertPrefix(tc.prefix, 1)
			if ok {
				t.Fatal("should not exist")
			}
			if _, ok := tree.GetPrefix(tc.prefix); !ok {
				t.Errorf("CIDR %s not found", tc.prefix)
			}
			if ok := tree.DeletePrefix(tc.prefix); !ok {
				t.Errorf("CIDR %s not deleted", tc.prefix)
			}
			if _, ok := tree.GetPrefix(tc.prefix); ok {
				t.Errorf("CIDR %s found", tc.prefix)
			}
		})
	}

}

func TestBasicIPv4(t *testing.T) {
	tree := NewIPTree[int]()
	// insert
	ipnet := netip.MustParsePrefix("192.168.0.0/24")
	ok := tree.InsertPrefix(ipnet, 1)
	if ok {
		t.Fatal("should not exist")
	}
	// check exist
	if _, ok := tree.GetPrefix(ipnet); !ok {
		t.Errorf("CIDR %s not found", ipnet)
	}

	// check does not exist
	ipnet2 := netip.MustParsePrefix("12.1.0.0/16")
	if _, ok := tree.GetPrefix(ipnet2); ok {
		t.Errorf("CIDR %s not expected", ipnet2)
	}

	// check insert existing prefix updates the value
	ok = tree.InsertPrefix(ipnet2, 2)
	if ok {
		t.Errorf("should not exist: %s", ipnet2)
	}

	ok = tree.InsertPrefix(ipnet2, 3)
	if !ok {
		t.Errorf("should be updated: %s", ipnet2)
	}

	if v, ok := tree.GetPrefix(ipnet2); !ok || v != 3 {
		t.Errorf("CIDR %s not expected", ipnet2)
	}

	// check longer prefix matching
	ipnet3 := netip.MustParsePrefix("12.1.0.2/32")
	lpm, _, ok := tree.LongestPrefixMatch(ipnet3)
	if !ok || lpm != ipnet2 {
		t.Errorf("expected %s got %s", ipnet2, lpm)
	}
}

func TestBasicIPv6(t *testing.T) {
	tree := NewIPTree[int]()
	// insert
	ipnet := netip.MustParsePrefix("2001:db8::/64")
	ok := tree.InsertPrefix(ipnet, 1)
	if ok {
		t.Fatal("should not exist")
	}
	// check exist
	if _, ok := tree.GetPrefix(ipnet); !ok {
		t.Errorf("CIDR %s not found", ipnet)
	}

	// check does not exist
	ipnet2 := netip.MustParsePrefix("2001:db8:1:3:4::/64")
	if _, ok := tree.GetPrefix(ipnet2); ok {
		t.Errorf("CIDR %s not expected", ipnet2)
	}

	// check insert existing prefix updates the value
	ok = tree.InsertPrefix(ipnet2, 2)
	if ok {
		t.Errorf("should not exist: %s", ipnet2)
	}

	ok = tree.InsertPrefix(ipnet2, 3)
	if !ok {
		t.Errorf("should be updated: %s", ipnet2)
	}

	if v, ok := tree.GetPrefix(ipnet2); !ok || v != 3 {
		t.Errorf("CIDR %s not expected", ipnet2)
	}

	// check longer prefix matching
	ipnet3 := netip.MustParsePrefix("2001:db8:1:3:4::/96")
	lpm, _, ok := tree.LongestPrefixMatch(ipnet3)
	if !ok || lpm != ipnet2 {
		t.Errorf("expected %s got %s", ipnet2, lpm)
	}
}

func TestInsertGetDelete100K(t *testing.T) {
	testCases := []struct {
		name string
		is6  bool
	}{
		{
			name: "ipv4",
		},
		{
			name: "ipv6",
			is6:  true,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			cidrs := generateRandomCIDRs(tc.is6, 100*1000)
			tree := NewIPTree[string]()

			for k := range cidrs {
				ok := tree.InsertPrefix(k, k.String())
				if ok {
					t.Errorf("error inserting: %v", k)
				}
			}

			if tree.Len(tc.is6) != len(cidrs) {
				t.Errorf("expected %d nodes on the tree, got %d", len(cidrs), tree.Len(tc.is6))
			}

			list := cidrs.UnsortedList()
			for _, k := range list {
				if v, ok := tree.GetPrefix(k); !ok {
					t.Errorf("CIDR %s not found", k)
					return
				} else if v != k.String() {
					t.Errorf("CIDR value %s not found", k)
					return
				}
				ok := tree.DeletePrefix(k)
				if !ok {
					t.Errorf("CIDR delete %s error", k)
				}
			}

			if tree.Len(tc.is6) != 0 {
				t.Errorf("No node expected on the tree, got: %d %v", tree.Len(tc.is6), cidrs)
			}
		})
	}
}

func Test_findAncestor(t *testing.T) {
	tests := []struct {
		name string
		a    netip.Prefix
		b    netip.Prefix
		want netip.Prefix
	}{
		{
			name: "ipv4 direct parent",
			a:    netip.MustParsePrefix("192.168.0.0/24"),
			b:    netip.MustParsePrefix("192.168.1.0/24"),
			want: netip.MustParsePrefix("192.168.0.0/23"),
		},
		{
			name: "ipv4 root parent ",
			a:    netip.MustParsePrefix("192.168.0.0/24"),
			b:    netip.MustParsePrefix("1.168.1.0/24"),
			want: netip.MustParsePrefix("0.0.0.0/0"),
		},
		{
			name: "ipv4 parent /1",
			a:    netip.MustParsePrefix("192.168.0.0/24"),
			b:    netip.MustParsePrefix("184.168.1.0/24"),
			want: netip.MustParsePrefix("128.0.0.0/1"),
		},
		{
			name: "ipv6 direct parent",
			a:    netip.MustParsePrefix("fd00:1:1:1::/64"),
			b:    netip.MustParsePrefix("fd00:1:1:2::/64"),
			want: netip.MustParsePrefix("fd00:1:1::/62"),
		},
		{
			name: "ipv6 root parent ",
			a:    netip.MustParsePrefix("fd00:1:1:1::/64"),
			b:    netip.MustParsePrefix("1:1:1:1::/64"),
			want: netip.MustParsePrefix("::/0"),
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := findAncestor(tt.a, tt.b); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("findAncestor() = %v, want %v", got, tt.want)
			}
		})
	}
}

func Test_getBitFromAddr(t *testing.T) {
	tests := []struct {
		name string
		ip   netip.Addr
		pos  int
		want int
	}{
		// 192.168.0.0
		// 11000000.10101000.00000000.00000001
		{
			name: "ipv4 first is a one",
			ip:   netip.MustParseAddr("192.168.0.0"),
			pos:  1,
			want: 1,
		},
		{
			name: "ipv4 middle is a zero",
			ip:   netip.MustParseAddr("192.168.0.0"),
			pos:  16,
			want: 0,
		},
		{
			name: "ipv4 middle is a one",
			ip:   netip.MustParseAddr("192.168.0.0"),
			pos:  13,
			want: 1,
		},
		{
			name: "ipv4 last is a zero",
			ip:   netip.MustParseAddr("192.168.0.0"),
			pos:  32,
			want: 0,
		},
		// 2001:db8::ff00:42:8329
		// 0010000000000001:0000110110111000:0000000000000000:0000000000000000:0000000000000000:1111111100000000:0000000001000010:1000001100101001
		{
			name: "ipv6 first is a zero",
			ip:   netip.MustParseAddr("2001:db8::ff00:42:8329"),
			pos:  1,
			want: 0,
		},
		{
			name: "ipv6 middle is a zero",
			ip:   netip.MustParseAddr("2001:db8::ff00:42:8329"),
			pos:  56,
			want: 0,
		},
		{
			name: "ipv6 middle is a one",
			ip:   netip.MustParseAddr("2001:db8::ff00:42:8329"),
			pos:  81,
			want: 1,
		},
		{
			name: "ipv6 last is a one",
			ip:   netip.MustParseAddr("2001:db8::ff00:42:8329"),
			pos:  128,
			want: 1,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := getBitFromAddr(tt.ip, tt.pos); got != tt.want {
				t.Errorf("getBitFromAddr() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestShortestPrefix(t *testing.T) {
	r := NewIPTree[int]()

	keys := []string{
		"10.0.0.0/8",
		"10.21.0.0/16",
		"10.221.0.0/16",
		"10.1.2.3/32",
		"10.1.2.0/24",
		"192.168.0.0/24",
		"192.168.0.0/16",
	}
	for _, k := range keys {
		ok := r.InsertPrefix(netip.MustParsePrefix(k), 0)
		if ok {
			t.Errorf("unexpected update on insert %s", k)
		}
	}
	if r.Len(false) != len(keys) {
		t.Fatalf("bad len: %v %v", r.Len(false), len(keys))
	}

	type exp struct {
		inp string
		out string
	}
	cases := []exp{
		{"192.168.0.3/32", "192.168.0.0/16"},
		{"10.1.2.4/21", "10.0.0.0/8"},
		{"192.168.0.0/16", "192.168.0.0/16"},
		{"192.168.0.0/32", "192.168.0.0/16"},
		{"10.1.2.3/32", "10.0.0.0/8"},
	}
	for _, test := range cases {
		m, _, ok := r.ShortestPrefixMatch(netip.MustParsePrefix(test.inp))
		if !ok {
			t.Fatalf("no match: %v", test)
		}
		if m != netip.MustParsePrefix(test.out) {
			t.Fatalf("mis-match: %v %v", m, test)
		}
	}

	// not match
	_, _, ok := r.ShortestPrefixMatch(netip.MustParsePrefix("0.0.0.0/0"))
	if ok {
		t.Fatalf("match unexpected for 0.0.0.0/0")
	}
}

func TestLongestPrefixMatch(t *testing.T) {
	r := NewIPTree[int]()

	keys := []string{
		"10.0.0.0/8",
		"10.21.0.0/16",
		"10.221.0.0/16",
		"10.1.2.3/32",
		"10.1.2.0/24",
		"192.168.0.0/24",
		"192.168.0.0/16",
	}
	for _, k := range keys {
		ok := r.InsertPrefix(netip.MustParsePrefix(k), 0)
		if ok {
			t.Errorf("unexpected update on insert %s", k)
		}
	}
	if r.Len(false) != len(keys) {
		t.Fatalf("bad len: %v %v", r.Len(false), len(keys))
	}

	type exp struct {
		inp string
		out string
	}
	cases := []exp{
		{"192.168.0.3/32", "192.168.0.0/24"},
		{"10.1.2.4/21", "10.0.0.0/8"},
		{"10.21.2.0/24", "10.21.0.0/16"},
		{"10.1.2.3/32", "10.1.2.3/32"},
	}
	for _, test := range cases {
		m, _, ok := r.LongestPrefixMatch(netip.MustParsePrefix(test.inp))
		if !ok {
			t.Fatalf("no match: %v", test)
		}
		if m != netip.MustParsePrefix(test.out) {
			t.Fatalf("mis-match: %v %v", m, test)
		}
	}
	// not match
	_, _, ok := r.LongestPrefixMatch(netip.MustParsePrefix("0.0.0.0/0"))
	if ok {
		t.Fatalf("match unexpected for 0.0.0.0/0")
	}
}

func TestTopLevelPrefixesV4(t *testing.T) {
	r := NewIPTree[string]()

	keys := []string{
		"10.0.0.0/8",
		"10.21.0.0/16",
		"10.221.0.0/16",
		"10.1.2.3/32",
		"10.1.2.0/24",
		"192.168.0.0/20",
		"192.168.1.0/24",
		"172.16.0.0/12",
		"172.21.23.0/24",
	}
	for _, k := range keys {
		ok := r.InsertPrefix(netip.MustParsePrefix(k), k)
		if ok {
			t.Errorf("unexpected update on insert %s", k)
		}
	}
	if r.Len(false) != len(keys) {
		t.Fatalf("bad len: %v %v", r.Len(false), len(keys))
	}

	expected := []string{
		"10.0.0.0/8",
		"192.168.0.0/20",
		"172.16.0.0/12",
	}
	parents := r.TopLevelPrefixes(false)
	if len(parents) != len(expected) {
		t.Fatalf("bad len: %v %v", len(parents), len(expected))
	}

	for _, k := range expected {
		v, ok := parents[netip.MustParsePrefix(k)]
		if !ok {
			t.Errorf("key %s not found", k)
		}
		if v != k {
			t.Errorf("value expected %s got %s", k, v)
		}
	}
}

func TestTopLevelPrefixesV6(t *testing.T) {
	r := NewIPTree[string]()

	keys := []string{
		"2001:db8:1:2:3::/64",
		"2001:db8::/64",
		"2001:db8:1:1:1::/64",
		"2001:db8:1:1:1::/112",
	}
	for _, k := range keys {
		ok := r.InsertPrefix(netip.MustParsePrefix(k), k)
		if ok {
			t.Errorf("unexpected update on insert %s", k)
		}
	}

	if r.Len(true) != len(keys) {
		t.Fatalf("bad len: %v %v", r.Len(true), len(keys))
	}

	expected := []string{
		"2001:db8::/64",
		"2001:db8:1:2:3::/64",
		"2001:db8:1:1:1::/64",
	}
	parents := r.TopLevelPrefixes(true)
	if len(parents) != len(expected) {
		t.Fatalf("bad len: %v %v", len(parents), len(expected))
	}

	for _, k := range expected {
		v, ok := parents[netip.MustParsePrefix(k)]
		if !ok {
			t.Errorf("key %s not found", k)
		}
		if v != k {
			t.Errorf("value expected %s got %s", k, v)
		}
	}
}

func TestWalkV4(t *testing.T) {
	r := NewIPTree[int]()

	keys := []string{
		"10.0.0.0/8",
		"10.1.0.0/16",
		"10.1.1.0/24",
		"10.1.1.32/26",
		"10.1.1.33/32",
	}
	for _, k := range keys {
		ok := r.InsertPrefix(netip.MustParsePrefix(k), 0)
		if ok {
			t.Errorf("unexpected update on insert %s", k)
		}
	}
	if r.Len(false) != len(keys) {
		t.Fatalf("bad len: %v %v", r.Len(false), len(keys))
	}

	// match exact prefix
	path := []string{}
	r.WalkPath(netip.MustParsePrefix("10.1.1.32/26"), func(k netip.Prefix, v int) bool {
		path = append(path, k.String())
		return false
	})
	if !cmp.Equal(path, keys[:4]) {
		t.Errorf("Walkpath expected %v got %v", keys[:4], path)
	}
	// not match on prefix
	path = []string{}
	r.WalkPath(netip.MustParsePrefix("10.1.1.33/26"), func(k netip.Prefix, v int) bool {
		path = append(path, k.String())
		return false
	})
	if !cmp.Equal(path, keys[:3]) {
		t.Errorf("Walkpath expected %v got %v", keys[:3], path)
	}
	// match exact prefix
	path = []string{}
	r.WalkPrefix(netip.MustParsePrefix("10.0.0.0/8"), func(k netip.Prefix, v int) bool {
		path = append(path, k.String())
		return false
	})
	if !cmp.Equal(path, keys) {
		t.Errorf("WalkPrefix expected %v got %v", keys, path)
	}
	// not match on prefix
	path = []string{}
	r.WalkPrefix(netip.MustParsePrefix("10.0.0.0/9"), func(k netip.Prefix, v int) bool {
		path = append(path, k.String())
		return false
	})
	if !cmp.Equal(path, keys[1:]) {
		t.Errorf("WalkPrefix expected %v got %v", keys[1:], path)
	}
}

func TestWalkV6(t *testing.T) {
	r := NewIPTree[int]()

	keys := []string{
		"2001:db8::/48",
		"2001:db8::/64",
		"2001:db8::/96",
		"2001:db8::/112",
		"2001:db8::/128",
	}
	for _, k := range keys {
		ok := r.InsertPrefix(netip.MustParsePrefix(k), 0)
		if ok {
			t.Errorf("unexpected update on insert %s", k)
		}
	}
	if r.Len(true) != len(keys) {
		t.Fatalf("bad len: %v %v", r.Len(false), len(keys))
	}

	// match exact prefix
	path := []string{}
	r.WalkPath(netip.MustParsePrefix("2001:db8::/112"), func(k netip.Prefix, v int) bool {
		path = append(path, k.String())
		return false
	})
	if !cmp.Equal(path, keys[:4]) {
		t.Errorf("Walkpath expected %v got %v", keys[:4], path)
	}
	// not match on prefix
	path = []string{}
	r.WalkPath(netip.MustParsePrefix("2001:db8::1/112"), func(k netip.Prefix, v int) bool {
		path = append(path, k.String())
		return false
	})
	if !cmp.Equal(path, keys[:3]) {
		t.Errorf("Walkpath expected %v got %v", keys[:3], path)
	}
	// match exact prefix
	path = []string{}
	r.WalkPrefix(netip.MustParsePrefix("2001:db8::/48"), func(k netip.Prefix, v int) bool {
		path = append(path, k.String())
		return false
	})
	if !cmp.Equal(path, keys) {
		t.Errorf("WalkPrefix expected %v got %v", keys, path)
	}
	// not match on prefix
	path = []string{}
	r.WalkPrefix(netip.MustParsePrefix("2001:db8::/49"), func(k netip.Prefix, v int) bool {
		path = append(path, k.String())
		return false
	})
	if !cmp.Equal(path, keys[1:]) {
		t.Errorf("WalkPrefix expected %v got %v", keys[1:], path)
	}
}

func TestGetHostIPPrefixMatches(t *testing.T) {
	r := NewIPTree[int]()

	keys := []string{
		"10.0.0.0/8",
		"10.21.0.0/16",
		"10.221.0.0/16",
		"10.1.2.3/32",
		"10.1.2.0/24",
		"192.168.0.0/24",
		"192.168.0.0/16",
		"2001:db8::/48",
		"2001:db8::/64",
		"2001:db8::/96",
	}
	for _, k := range keys {
		ok := r.InsertPrefix(netip.MustParsePrefix(k), 0)
		if ok {
			t.Errorf("unexpected update on insert %s", k)
		}
	}

	type exp struct {
		inp string
		out []string
	}
	cases := []exp{
		{"192.168.0.3", []string{"192.168.0.0/24", "192.168.0.0/16"}},
		{"10.1.2.4", []string{"10.1.2.0/24", "10.0.0.0/8"}},
		{"10.1.2.0", []string{"10.0.0.0/8"}},
		{"10.1.2.255", []string{"10.0.0.0/8"}},
		{"192.168.0.0", []string{}},
		{"192.168.1.0", []string{"192.168.0.0/16"}},
		{"10.1.2.255", []string{"10.0.0.0/8"}},
		{"2001:db8::1", []string{"2001:db8::/96", "2001:db8::/64", "2001:db8::/48"}},
		{"2001:db8::", []string{}},
		{"2001:db8::ffff:ffff:ffff:ffff", []string{"2001:db8::/64", "2001:db8::/48"}},
	}
	for _, test := range cases {
		m := r.GetHostIPPrefixMatches(netip.MustParseAddr(test.inp))
		in := []netip.Prefix{}
		for k := range m {
			in = append(in, k)
		}
		out := []netip.Prefix{}
		for _, s := range test.out {
			out = append(out, netip.MustParsePrefix(s))
		}

		// sort by prefix bits to avoid flakes
		sort.Slice(in, func(i, j int) bool { return in[i].Bits() < in[j].Bits() })
		sort.Slice(out, func(i, j int) bool { return out[i].Bits() < out[j].Bits() })
		if !reflect.DeepEqual(in, out) {
			t.Fatalf("mis-match: %v %v", in, out)
		}
	}

	// not match
	_, _, ok := r.ShortestPrefixMatch(netip.MustParsePrefix("0.0.0.0/0"))
	if ok {
		t.Fatalf("match unexpected for 0.0.0.0/0")
	}
}

func Test_prefixContainIP(t *testing.T) {
	tests := []struct {
		name   string
		prefix netip.Prefix
		ip     netip.Addr
		want   bool
	}{
		{
			name:   "IPv4 contains",
			prefix: netip.MustParsePrefix("192.168.0.0/24"),
			ip:     netip.MustParseAddr("192.168.0.1"),
			want:   true,
		},
		{
			name:   "IPv4 network address",
			prefix: netip.MustParsePrefix("192.168.0.0/24"),
			ip:     netip.MustParseAddr("192.168.0.0"),
		},
		{
			name:   "IPv4 broadcast address",
			prefix: netip.MustParsePrefix("192.168.0.0/24"),
			ip:     netip.MustParseAddr("192.168.0.255"),
		},
		{
			name:   "IPv4 does not contain",
			prefix: netip.MustParsePrefix("192.168.0.0/24"),
			ip:     netip.MustParseAddr("192.168.1.2"),
		},
		{
			name:   "IPv6 contains",
			prefix: netip.MustParsePrefix("2001:db2::/96"),
			ip:     netip.MustParseAddr("2001:db2::1"),
			want:   true,
		},
		{
			name:   "IPv6 network address",
			prefix: netip.MustParsePrefix("2001:db2::/96"),
			ip:     netip.MustParseAddr("2001:db2::"),
		},
		{
			name:   "IPv6 broadcast address",
			prefix: netip.MustParsePrefix("2001:db2::/96"),
			ip:     netip.MustParseAddr("2001:db2::ffff:ffff"),
			want:   true,
		},
		{
			name:   "IPv6 does not contain",
			prefix: netip.MustParsePrefix("2001:db2::/96"),
			ip:     netip.MustParseAddr("2001:db2:1:2:3::1"),
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := prefixContainIP(tt.prefix, tt.ip); got != tt.want {
				t.Errorf("prefixContainIP() = %v, want %v", got, tt.want)
			}
		})
	}
}

func BenchmarkInsertUpdate(b *testing.B) {
	r := NewIPTree[bool]()
	ipList := generateRandomCIDRs(true, 20000).UnsortedList()
	for _, ip := range ipList {
		r.InsertPrefix(ip, true)
	}

	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		r.InsertPrefix(ipList[n%len(ipList)], true)
	}
}

func generateRandomCIDRs(is6 bool, number int) sets.Set[netip.Prefix] {
	n := 4
	if is6 {
		n = 16
	}
	cidrs := sets.Set[netip.Prefix]{}
	rand.New(rand.NewSource(time.Now().UnixNano()))
	for i := 0; i < number; i++ {
		bytes := make([]byte, n)
		for i := 0; i < n; i++ {
			bytes[i] = uint8(rand.Intn(255))
		}

		ip, ok := netip.AddrFromSlice(bytes)
		if !ok {
			continue
		}

		bits := rand.Intn(n * 8)
		prefix := netip.PrefixFrom(ip, bits).Masked()
		if prefix.IsValid() {
			cidrs.Insert(prefix)
		}
	}
	return cidrs
}
