package nettrie

import (
	"maps"
	"net/netip"
	"reflect"
	"slices"
	"testing"
)

func TestValueTrie_IPv4(t *testing.T) {
	trie := NewValue[string]()

	routes := map[string]string{
		"0.0.0.0/0":        "Default",
		"10.0.0.0/8":       "Private A",
		"192.168.0.0/16":   "Private B",
		"192.168.1.0/24":   "LAN",
		"192.168.1.128/25": "Subnet",
		"192.168.1.200/32": "Host",
	}

	for s, v := range routes {
		trie.Insert(netip.MustParsePrefix(s), v)
	}

	testCases := []struct {
		name        string
		lookupAddr  string
		expectedVal string
		expectedOK  bool
	}{
		{"Exact Host Match", "192.168.1.200", "Host", true},
		{"Subnet Match", "192.168.1.201", "Subnet", true},
		{"LAN Match", "192.168.1.50", "LAN", true},
		{"Private B Match", "192.168.255.255", "Private B", true},
		{"Private A Match", "10.255.255.255", "Private A", true},
		{"Default Match", "8.8.8.8", "Default", true},
		{"No Match", "203.0.113.1", "Default", true}, // Falls back to default
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			addr := netip.MustParseAddr(tc.lookupAddr)
			val, ok := trie.Lookup(addr)

			if ok != tc.expectedOK {
				t.Errorf("expected ok=%v, got ok=%v", tc.expectedOK, ok)
			}
			if val != tc.expectedVal {
				t.Errorf("expected value=%q, got %q", tc.expectedVal, val)
			}
		})
	}
}

func TestValueTrie_IPv6(t *testing.T) {
	trie := NewValue[string]()

	routes := map[string]string{
		"::/0":                   "Default",
		"2001:db8::/32":          "Global Unicast",
		"2001:db8:acad::/48":     "Academic",
		"2001:db8:acad:1::/64":   "CS Department",
		"2001:db8:acad:1::1/128": "Host Route",
	}

	for s, v := range routes {
		trie.Insert(netip.MustParsePrefix(s), v)
	}

	testCases := []struct {
		name        string
		lookupAddr  string
		expectedVal string
		expectedOK  bool
	}{
		{"Exact Host Match", "2001:db8:acad:1::1", "Host Route", true},
		{"CS Dept Match", "2001:db8:acad:1::2", "CS Department", true},
		{"Academic Match", "2001:db8:acad:2::1", "Academic", true},
		{"Global Unicast Match", "2001:db8:cafe::1", "Global Unicast", true},
		{"Default Match", "2606:4700:4700::1111", "Default", true},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			addr := netip.MustParseAddr(tc.lookupAddr)
			val, ok := trie.Lookup(addr)

			if ok != tc.expectedOK {
				t.Errorf("expected ok=%v, got ok=%v", tc.expectedOK, ok)
			}
			if val != tc.expectedVal {
				t.Errorf("expected value=%q, got %q", tc.expectedVal, val)
			}
		})
	}
}

func TestValueTrie_UpdateValue(t *testing.T) {
	trie := NewValue[string]()
	prefix := netip.MustParsePrefix("192.168.1.0/24")
	addr := netip.MustParseAddr("192.168.1.1")

	trie.Insert(prefix, "Initial Value")
	val, _ := trie.Lookup(addr)
	if val != "Initial Value" {
		t.Fatalf("expected initial value, got %q", val)
	}

	trie.Insert(prefix, "Updated Value")
	val, _ = trie.Lookup(addr)
	if val != "Updated Value" {
		t.Fatalf("expected updated value, got %q", val)
	}
}

func TestValueTrie_Delete(t *testing.T) {
	trie := NewValue[string]()

	trie.Insert(netip.MustParsePrefix("10.0.0.0/8"), "A")
	trie.Insert(netip.MustParsePrefix("10.0.0.0/24"), "B")
	trie.Insert(netip.MustParsePrefix("10.0.1.0/24"), "C")

	t.Run("Delete Non-Existent", func(t *testing.T) {
		if trie.Delete(netip.MustParsePrefix("1.2.3.4/32")) {
			t.Error("deleted a non-existent prefix")
		}
	})

	t.Run("Delete Leaf Node", func(t *testing.T) {
		// Delete 10.0.0.0/24, which is a leaf.
		if !trie.Delete(netip.MustParsePrefix("10.0.0.0/24")) {
			t.Fatal("failed to delete 10.0.0.0/24")
		}
		// Lookup should now resolve to the parent 10.0.0.0/8
		val, ok := trie.Lookup(netip.MustParseAddr("10.0.0.1"))
		if !ok || val != "A" {
			t.Errorf("lookup failed after delete, got val=%q ok=%v, want val=\"A\" ok=true", val, ok)
		}
	})

	t.Run("Delete and Merge", func(t *testing.T) {
		// Insert a new prefix that causes a split
		trie.Insert(netip.MustParsePrefix("10.0.0.0/8"), "A-updated")

		// We now have 10.0.0.0/8 and 10.0.1.0/24. Deleting 10.0.1.0/24 should
		// cause the split node to be removed and merged back into 10.0.0.0/8.
		if !trie.Delete(netip.MustParsePrefix("10.0.1.0/24")) {
			t.Fatal("failed to delete 10.0.1.0/24")
		}

		// A lookup for 10.0.1.1 should now match 10.0.0.0/8
		val, ok := trie.Lookup(netip.MustParseAddr("10.0.1.1"))
		if !ok || val != "A-updated" {
			t.Errorf("lookup failed after merge, got val=%q ok=%v, want val=\"A-updated\" ok=true", val, ok)
		}
	})

	t.Run("Delete Prefix Of Another", func(t *testing.T) {
		// Delete 10.0.0.0/8
		trie.Insert(netip.MustParsePrefix("10.0.0.0/8"), "A")
		trie.Insert(netip.MustParsePrefix("10.1.0.0/16"), "D")

		if !trie.Delete(netip.MustParsePrefix("10.0.0.0/8")) {
			t.Fatal("failed to delete 10.0.0.0/8")
		}

		// Lookup for 10.0.0.1 should now fail (no default route)
		_, ok := trie.Lookup(netip.MustParseAddr("10.0.0.1"))
		if ok {
			t.Error("lookup for 10.0.0.1 should have failed")
		}

		// Lookup for 10.1.0.1 should still succeed
		val, ok := trie.Lookup(netip.MustParseAddr("10.1.0.1"))
		if !ok || val != "D" {
			t.Error("lookup for 10.1.0.1 failed unexpectedly")
		}
	})
}

func TestValueTrie_Contains(t *testing.T) {
	trie := NewValue[string]()

	prefixes := []string{
		"10.0.0.0/8",
		"192.168.1.0/24",
		"192.168.1.200/32",
		"2001:db8::/32",
		"2001:db8:acad:1::1/128",
	}

	for _, s := range prefixes {
		trie.Insert(netip.MustParsePrefix(s), "present")
	}

	t.Run("ContainsPrefix", func(t *testing.T) {
		testCases := []struct {
			prefix string
			want   bool
		}{
			{"10.0.0.0/8", true},
			{"192.168.1.200/32", true},
			{"2001:db8::/32", true},
			{"2001:db8:acad:1::1/128", true},
			{"10.0.0.0/9", false},     // shorter parent, but not exact
			{"192.168.1.0/25", false}, // non-existent child
			{"172.16.0.0/12", false},  // completely different prefix
			{"2001:db8::/48", false},  // non-existent child
		}

		for _, tc := range testCases {
			p := netip.MustParsePrefix(tc.prefix)
			if got := trie.ContainsPrefix(p); got != tc.want {
				t.Errorf("ContainsPrefix(%q) = %v, want %v", tc.prefix, got, tc.want)
			}
		}
	})

	t.Run("Contains", func(t *testing.T) {
		testCases := []struct {
			addr string
			want bool
		}{
			{"192.168.1.200", true},
			{"2001:db8:acad:1::1", true},
			{"192.168.1.201", false}, // In /24 range, but not a /32 host route
			{"10.0.0.1", false},      // In /8 range, but not a /32 host route
		}

		for _, tc := range testCases {
			a := netip.MustParseAddr(tc.addr)
			if got := trie.Contains(a); got != tc.want {
				t.Errorf("Contains(%q) = %v, want %v", tc.addr, got, tc.want)
			}
		}
	})
}

func TestValueTrie_Walk(t *testing.T) {
	trie := NewValue[string]()
	prefixes := map[string]string{
		"10.0.0.0/8":         "A",
		"192.168.1.0/24":     "B",
		"2001:db8::/32":      "C",
		"172.16.0.0/12":      "D",
		"2001:db8:acad::/48": "E",
	}
	for s, v := range prefixes {
		trie.Insert(netip.MustParsePrefix(s), v)
	}

	t.Run("Walk all prefixes", func(t *testing.T) {
		walked := make(map[string]string)
		trie.Walk(func(p netip.Prefix, v string) bool {
			walked[p.String()] = v
			return true
		})

		if !maps.Equal(walked, prefixes) {
			t.Errorf("walked prefixes mismatch:\nexpected: %v\ngot:      %v", prefixes, walked)
		}
	})

	t.Run("Stop walk early", func(t *testing.T) {
		count := 0
		trie.Walk(func(p netip.Prefix, v string) bool {
			count++
			return count < 3 // Stop after visiting 3 prefixes
		})

		if count != 3 {
			t.Errorf("expected walk to stop after 3 prefixes, but it visited %d", count)
		}
	})

	t.Run("Stop walk between families", func(t *testing.T) {
		stopTrie := NewValue[int]()
		stopTrie.Insert(netip.MustParsePrefix("10.0.0.0/8"), 1)
		stopTrie.Insert(netip.MustParsePrefix("2001:db8::/32"), 2)

		count := 0
		stopTrie.Walk(func(p netip.Prefix, v int) bool {
			count++
			return false
		})

		if count != 1 {
			t.Errorf("expected walk to stop after 1 prefix, but it visited %d", count)
		}
	})
}

func TestValueTrie_Merge(t *testing.T) {
	trieA := NewValue[string]()
	trieA.Insert(netip.MustParsePrefix("10.0.0.0/8"), "net10")
	trieA.Insert(netip.MustParsePrefix("192.168.1.0/24"), "lan_A")
	trieA.Insert(netip.MustParsePrefix("2001:db8::/32"), "v6_A")

	trieB := NewValue[string]()
	trieB.Insert(netip.MustParsePrefix("10.1.0.0/16"), "net10_subnet")
	trieB.Insert(netip.MustParsePrefix("192.168.1.0/24"), "lan_B_override") // Overlap
	trieB.Insert(netip.MustParsePrefix("172.16.0.0/12"), "corp")
	trieB.Insert(netip.MustParsePrefix("2001:db8:acad::/48"), "v6_B")

	trieA.Merge(trieB)

	expected := map[string]string{
		"10.0.0.0/8":         "net10",
		"192.168.1.0/24":     "lan_B_override", // Value should be from trieB
		"2001:db8::/32":      "v6_A",
		"10.1.0.0/16":        "net10_subnet",
		"172.16.0.0/12":      "corp",
		"2001:db8:acad::/48": "v6_B",
	}

	actual := make(map[string]string)
	trieA.Walk(func(p netip.Prefix, v string) bool {
		actual[p.String()] = v
		return true
	})

	if !reflect.DeepEqual(expected, actual) {
		t.Errorf("Merge result incorrect.\nExpected: %v\nGot: %v", expected, actual)
	}

	// Verify trieB is unchanged by collecting its prefixes
	bPrefixes := []string{}
	trieB.Walk(func(p netip.Prefix, v string) bool {
		bPrefixes = append(bPrefixes, p.String())
		return true
	})
	expectedBPrefixes := []string{"10.1.0.0/16", "172.16.0.0/12", "192.168.1.0/24", "2001:db8:acad::/48"}
	slices.Sort(bPrefixes)
	slices.Sort(expectedBPrefixes)
	if !reflect.DeepEqual(bPrefixes, expectedBPrefixes) {
		t.Errorf("trieB was modified during merge.\nExpected: %v\nGot: %v", expectedBPrefixes, bPrefixes)
	}
}