package routetable_test

import (
	"net/netip"
	"reflect"
	"testing"

	"git.eeqj.de/sneak/routewatch/internal/routetable"
)

// announceAll has each listed peer announce prefix with origin AS asn, the
// peer at index i i*10 seconds after start.
func announceAll(
	table *routetable.Table, prefix string, asn, start int, peers ...string,
) {
	const secondsBetweenPeers = 10

	for i, peer := range peers {
		table.Announce(netip.MustParsePrefix(prefix), asn,
			routeOf(peer, asn, start+i*secondsBetweenPeers))
	}
}

func TestLookup(t *testing.T) {
	t.Parallel()

	table := routetable.New(routetable.Params{})
	announceAll(table, "10.0.0.0/8", asA, 10, peerA)
	// Three peers with two origins: asB has more of them.
	announceAll(table, "10.1.0.0/16", asA, 10, peerA)
	announceAll(table, "10.1.0.0/16", asB, 20, peerB, peerC)
	// Two origins with one peer each: the lower AS number wins.
	announceAll(table, prefixV4, asB, 20, peerB)
	announceAll(table, prefixV4, asA, 10, peerA)
	// An origin whose kept routes were withdrawn: no newest time.
	announceAll(table, "10.9.0.0/16", asA, 10, peerA, peerB, peerC)
	table.Withdraw(netip.MustParsePrefix("10.9.0.0/16"), peerB)
	table.Withdraw(netip.MustParsePrefix("10.9.0.0/16"), peerC)
	announceAll(table, "2001:db8::/32", asA, 10, peerA)
	announceAll(table, prefixV6, asB, 20, peerB)
	announceAll(table, "3000::/16", asA, 10, peerA)
	// Shorter than the lookup uses.
	announceAll(table, "192.0.0.0/7", asA, 10, peerA)
	announceAll(table, "2000::/15", asA, 10, peerA)

	tests := []struct {
		addr string
		// want is the summary's prefix, or empty when none is found.
		want      string
		originASN int
		peers     int
		newest    int // in seconds after the start; 0 for none
	}{
		{"10.1.2.3", prefixV4, asA, 2, 20},
		{"::ffff:10.1.2.3", prefixV4, asA, 2, 20},
		{"10.1.3.4", "10.1.0.0/16", asB, 3, 30},
		{"10.2.0.1", "10.0.0.0/8", asA, 1, 10},
		{"10.9.0.1", "10.9.0.0/16", asA, 1, 0},
		{"192.0.2.1", "", 0, 0, 0},
		{"2001:db8:1::5", prefixV6, asB, 1, 20},
		{"2001:db8:2::1", "2001:db8::/32", asA, 1, 10},
		{"3000::1", "3000::/16", asA, 1, 10},
		{"2001:db9::1", "", 0, 0, 0},
	}

	for _, tc := range tests {
		t.Run(tc.addr, func(t *testing.T) {
			t.Parallel()

			got, found := table.Lookup(netip.MustParseAddr(tc.addr))
			if tc.want == "" {
				if found {
					t.Errorf("Lookup(%s) = %+v, want none", tc.addr, got)
				}

				return
			}

			want := routetable.PrefixSummary{
				Prefix:    netip.MustParsePrefix(tc.want),
				OriginASN: tc.originASN,
				Peers:     tc.peers,
			}
			if tc.newest != 0 {
				want.Newest = timeAt(tc.newest)
			}

			if !found || !reflect.DeepEqual(got, want) {
				t.Errorf("Lookup(%s) = %+v, %v; want %+v", tc.addr, got, found, want)
			}
		})
	}
}

func TestOriginatedPrefixes(t *testing.T) {
	t.Parallel()

	table := routetable.New(routetable.Params{})
	announceAll(table, "2001:db8::/32", asA, 30, peerA)
	announceAll(table, prefixV4, asA, 10, peerA)
	// asB's newer route for the same prefix does not count for asA.
	announceAll(table, prefixV4, asB, 40, peerB)
	announceAll(table, "10.1.0.0/16", asA, 10, peerA, peerB)
	announceAll(table, "10.9.0.0/16", asA, 10, peerA, peerB, peerC)
	table.Withdraw(netip.MustParsePrefix("10.9.0.0/16"), peerB)
	table.Withdraw(netip.MustParsePrefix("10.9.0.0/16"), peerC)

	want := []routetable.OriginatedPrefix{
		{Prefix: netip.MustParsePrefix("10.1.0.0/16"), Newest: timeAt(20)},
		{Prefix: netip.MustParsePrefix(prefixV4), Newest: timeAt(10)},
		{Prefix: netip.MustParsePrefix("10.9.0.0/16")},
		{Prefix: netip.MustParsePrefix("2001:db8::/32"), Newest: timeAt(30)},
	}

	got := table.OriginatedPrefixes(asA)
	if !reflect.DeepEqual(got, want) {
		t.Errorf("OriginatedPrefixes(%d) = %+v, want %+v", asA, got, want)
	}

	got = table.OriginatedPrefixes(asC)
	if len(got) != 0 {
		t.Errorf("OriginatedPrefixes(%d) = %+v, want none", asC, got)
	}
}

func TestSamplePrefixes(t *testing.T) {
	t.Parallel()

	table := routetable.New(routetable.Params{})

	announced := []string{
		"10.0.0.0/24", "10.0.1.0/24", "10.0.2.0/24", "10.0.3.0/24",
		"10.0.4.0/24", "10.0.0.0/16",
		"2001:db8::/48", "2001:db8:1::/48", "2001:db8:2::/48",
	}
	for _, prefix := range announced {
		announceAll(table, prefix, asA, 10, peerA)
	}

	tests := []struct {
		name                         string
		ipVersion, maskLength, limit int
		want                         int
	}{
		{"fewer than there are", 4, 24, 3, 3},
		{"more than there are", 4, 24, 10, 5},
		{"IPv6", 6, 48, 10, 3},
		{"a mask length with none", 4, 20, 10, 0},
		{"no such IP version", 5, 24, 10, 0},
		{"a mask length too long", 4, 33, 10, 0},
		{"a negative mask length", 4, -1, 10, 0},
		{"a limit of zero", 4, 24, 0, 0},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			sample := table.SamplePrefixes(tc.ipVersion, tc.maskLength, tc.limit)
			if len(sample) != tc.want {
				t.Fatalf("SamplePrefixes(%d, %d, %d) returned %d prefixes, want %d",
					tc.ipVersion, tc.maskLength, tc.limit, len(sample), tc.want)
			}

			seen := make(map[netip.Prefix]bool)

			for _, got := range sample {
				want := routetable.PrefixSummary{
					Prefix: got.Prefix, OriginASN: asA, Peers: 1, Newest: timeAt(10),
				}
				if seen[got.Prefix] || got.Prefix.Bits() != tc.maskLength ||
					!reflect.DeepEqual(got, want) {
					t.Errorf("SamplePrefixes(%d, %d, %d) returned %+v",
						tc.ipVersion, tc.maskLength, tc.limit, got)
				}

				seen[got.Prefix] = true
			}
		})
	}
}

func TestCounts(t *testing.T) {
	t.Parallel()

	table := routetable.New(routetable.Params{})
	checkCounts(t, table, routetable.Counts{
		IPv4PrefixesByLength: byLength(ipv4Bits, nil),
		IPv6PrefixesByLength: byLength(ipv6Bits, nil),
	})

	announceAll(table, prefixV4, asA, 10, peerA, peerB)
	announceAll(table, "10.2.0.0/16", asA, 15, peerA)
	announceAll(table, "2001:db8::/32", asA, 5, peerA)
	announceAll(table, "2001:db8::/32", asA, 25, peerC)
	checkCounts(t, table, routetable.Counts{
		Pairs:                3,
		Peers:                3,
		IPv4Prefixes:         2,
		IPv6Prefixes:         1,
		IPv4PrefixesByLength: byLength(ipv4Bits, map[int]int{16: 1, 24: 1}),
		IPv6PrefixesByLength: byLength(ipv6Bits, map[int]int{32: 1}),
		IPv4Routes:           3,
		IPv6Routes:           2,
		OldestKeptRoute:      timeAt(5),
		NewestKeptRoute:      timeAt(25),
	})

	// peerC's route drops peerA's from the kept ones, and peerA leaves the
	// IPv6 prefix.
	announceAll(table, prefixV4, asA, 30, peerC)
	table.Withdraw(netip.MustParsePrefix("2001:db8::/32"), peerA)
	checkCounts(t, table, routetable.Counts{
		Pairs:                3,
		Peers:                3,
		IPv4Prefixes:         2,
		IPv6Prefixes:         1,
		IPv4PrefixesByLength: byLength(ipv4Bits, map[int]int{16: 1, 24: 1}),
		IPv6PrefixesByLength: byLength(ipv6Bits, map[int]int{32: 1}),
		IPv4Routes:           4,
		IPv6Routes:           1,
		OldestKeptRoute:      timeAt(15),
		NewestKeptRoute:      timeAt(30),
	})

	// peerC was the last peer of the IPv6 prefix. It keeps its number.
	table.PeerDown(peerC)
	checkCounts(t, table, routetable.Counts{
		Pairs:                2,
		Peers:                3,
		IPv4Prefixes:         2,
		IPv4PrefixesByLength: byLength(ipv4Bits, map[int]int{16: 1, 24: 1}),
		IPv6PrefixesByLength: byLength(ipv6Bits, nil),
		IPv4Routes:           3,
		OldestKeptRoute:      timeAt(15),
		NewestKeptRoute:      timeAt(20),
	})
}

// The address lengths of the two IP versions.
const (
	ipv4Bits = 32
	ipv6Bits = 128
)

// byLength returns the prefix counts by mask length, for addresses of bits
// bits, with counts[n] prefixes of mask length n and none of any other.
func byLength(bits int, counts map[int]int) []int {
	all := make([]int, bits+1)
	for n, count := range counts {
		all[n] = count
	}

	return all
}

// checkCounts checks that the table's counts are want.
func checkCounts(t *testing.T, table *routetable.Table, want routetable.Counts) {
	t.Helper()

	got := table.Counts()
	if !reflect.DeepEqual(got, want) {
		t.Errorf("Counts() = %+v, want %+v", got, want)
	}
}