package routetable_test

import (
	"net/netip"
	"reflect"
	"slices"
	"sync"
	"testing"
	"time"

	"github.com/bits-and-blooms/bitset"

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

const (
	peerA = "192.0.2.1"
	peerB = "192.0.2.2"
	peerC = "2001:db8::3"
	peerD = "2001:db8::4"

	asA = 64500
	asB = 64501
	asC = 64502
	// transitAS is the first AS of every test path, the peer's own.
	transitAS = 64496

	prefixV4 = "10.1.2.0/24"
	prefixV6 = "2001:db8:1::/48"
)

// at returns the time second seconds after the start of every test.
func at(second int) time.Time {
	return time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC).
		Add(time.Duration(second) * time.Second)
}

// timeAt returns at(second) for a field that may be nil.
func timeAt(second int) *time.Time {
	t := at(second)

	return &t
}

// routeOf returns the route peer announces with origin AS asn second seconds
// after the start.
func routeOf(peer string, asn, second int) routetable.Route {
	return routetable.Route{
		Peer: peer, ASPath: []int{transitAS, asn}, Time: at(second),
	}
}

// step is one update a test applies to a table for one prefix.
type step func(table *routetable.Table, prefix netip.Prefix)

// announce is the step of peer announcing prefix with origin AS asn second
// seconds after the start.
func announce(peer string, asn, second int) step {
	return func(table *routetable.Table, prefix netip.Prefix) {
		table.Announce(prefix, asn, routeOf(peer, asn, second))
	}
}

// withdraw is the step of peer withdrawing prefix.
func withdraw(peer string) step {
	return func(table *routetable.Table, prefix netip.Prefix) {
		table.Withdraw(prefix, peer)
	}
}

// peerDown is the step of peer's session going down.
func peerDown(peer string) step {
	return func(table *routetable.Table, _ netip.Prefix) {
		table.PeerDown(peer)
	}
}

// rule is one case of the rule that decides what the table holds for a prefix:
// the steps applied to an empty table, and the prefix's origins after them.
// The peers are numbered in the order they first announce, which is the order
// Origin.Peers lists them in.
type rule struct {
	name  string
	steps []step
	want  []routetable.Origin
}

// checkRules runs each case for an IPv4 and an IPv6 prefix.
func checkRules(t *testing.T, rules []rule) {
	t.Helper()

	for _, tc := range rules {
		for _, text := range []string{prefixV4, prefixV6} {
			t.Run(tc.name+" "+text, func(t *testing.T) {
				t.Parallel()

				prefix := netip.MustParsePrefix(text)
				table := routetable.New(routetable.Params{})

				for _, apply := range tc.steps {
					apply(table, prefix)
				}

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

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

	checkRules(t, []rule{
		{
			name:  "an announcement is kept",
			steps: []step{announce(peerA, asA, 10)},
			want: []routetable.Origin{{ASN: asA, Peers: []string{peerA},
				Routes: []routetable.Route{routeOf(peerA, asA, 10)}}},
		},
		{
			name: "a lagging peer's older route is not kept",
			steps: []step{announce(peerA, asA, 20), announce(peerB, asA, 30),
				announce(peerC, asA, 10)},
			want: []routetable.Origin{{ASN: asA,
				Peers: []string{peerA, peerB, peerC}, Routes: []routetable.Route{
					routeOf(peerB, asA, 30), routeOf(peerA, asA, 20),
				}}},
		},
		{
			name: "a third newer route drops the older kept one, whose peer " +
				"keeps its bit",
			steps: []step{announce(peerA, asA, 10), announce(peerB, asA, 20),
				announce(peerC, asA, 30)},
			want: []routetable.Origin{{ASN: asA,
				Peers: []string{peerA, peerB, peerC}, Routes: []routetable.Route{
					routeOf(peerC, asA, 30), routeOf(peerB, asA, 20),
				}}},
		},
		{
			name: "a route as old as the older kept one replaces it",
			steps: []step{announce(peerA, asA, 10), announce(peerB, asA, 20),
				announce(peerC, asA, 10)},
			want: []routetable.Origin{{ASN: asA,
				Peers: []string{peerA, peerB, peerC}, Routes: []routetable.Route{
					routeOf(peerB, asA, 20), routeOf(peerC, asA, 10),
				}}},
		},
		{
			name: "a peer's new route replaces its kept one",
			steps: []step{announce(peerA, asA, 10), announce(peerB, asA, 20),
				announce(peerA, asA, 30)},
			want: []routetable.Origin{{ASN: asA, Peers: []string{peerA, peerB},
				Routes: []routetable.Route{
					routeOf(peerA, asA, 30), routeOf(peerB, asA, 20),
				}}},
		},
		{
			name: "a peer's new route replaces its kept one even when older",
			steps: []step{announce(peerA, asA, 10), announce(peerB, asA, 20),
				announce(peerB, asA, 5)},
			want: []routetable.Origin{{ASN: asA, Peers: []string{peerA, peerB},
				Routes: []routetable.Route{
					routeOf(peerA, asA, 10), routeOf(peerB, asA, 5),
				}}},
		},
	})
}

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

	checkRules(t, []rule{
		{
			name: "a peer moving to a new origin leaves the old one",
			steps: []step{announce(peerA, asA, 10), announce(peerB, asA, 20),
				announce(peerA, asB, 30)},
			want: []routetable.Origin{
				{ASN: asA, Peers: []string{peerB},
					Routes: []routetable.Route{routeOf(peerB, asA, 20)}},
				{ASN: asB, Peers: []string{peerA},
					Routes: []routetable.Route{routeOf(peerA, asB, 30)}},
			},
		},
		{
			name: "a peer whose route is not kept moves to a new origin",
			steps: []step{announce(peerA, asA, 10), announce(peerB, asA, 20),
				announce(peerC, asA, 30), announce(peerA, asB, 40)},
			want: []routetable.Origin{
				{ASN: asA, Peers: []string{peerB, peerC},
					Routes: []routetable.Route{
						routeOf(peerC, asA, 30), routeOf(peerB, asA, 20),
					}},
				{ASN: asB, Peers: []string{peerA},
					Routes: []routetable.Route{routeOf(peerA, asB, 40)}},
			},
		},
		{
			name: "origins are listed in AS number order",
			steps: []step{announce(peerA, asC, 10), announce(peerB, asA, 20),
				announce(peerC, asB, 30)},
			want: []routetable.Origin{
				{ASN: asA, Peers: []string{peerB},
					Routes: []routetable.Route{routeOf(peerB, asA, 20)}},
				{ASN: asB, Peers: []string{peerC},
					Routes: []routetable.Route{routeOf(peerC, asB, 30)}},
				{ASN: asC, Peers: []string{peerA},
					Routes: []routetable.Route{routeOf(peerA, asC, 10)}},
			},
		},
	})
}

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

	// threePeers is peerA, peerB and peerC announcing in turn; the routes of
	// peerC and peerB are kept.
	threePeers := []step{announce(peerA, asA, 10), announce(peerB, asA, 20),
		announce(peerC, asA, 30)}

	checkRules(t, []rule{
		{
			name:  "a withdrawal drops the peer's kept route",
			steps: append(slices.Clone(threePeers), withdraw(peerC)),
			want: []routetable.Origin{{ASN: asA, Peers: []string{peerA, peerB},
				Routes: []routetable.Route{routeOf(peerB, asA, 20)}}},
		},
		{
			name:  "a withdrawal of a route that is not kept leaves the kept ones",
			steps: append(slices.Clone(threePeers), withdraw(peerA)),
			want: []routetable.Origin{{ASN: asA, Peers: []string{peerB, peerC},
				Routes: []routetable.Route{
					routeOf(peerC, asA, 30), routeOf(peerB, asA, 20),
				}}},
		},
		{
			name: "an origin whose kept routes are withdrawn keeps its other " +
				"peers",
			steps: append(slices.Clone(threePeers), withdraw(peerB),
				withdraw(peerC)),
			want: []routetable.Origin{{ASN: asA, Peers: []string{peerA}}},
		},
		{
			name: "an origin without kept routes keeps the next update",
			steps: append(slices.Clone(threePeers), withdraw(peerB),
				withdraw(peerC), announce(peerA, asA, 5)),
			want: []routetable.Origin{{ASN: asA, Peers: []string{peerA},
				Routes: []routetable.Route{routeOf(peerA, asA, 5)}}},
		},
		{
			name: "an origin no peer carries is removed",
			steps: []step{announce(peerA, asA, 10), announce(peerB, asB, 20),
				withdraw(peerA)},
			want: []routetable.Origin{{ASN: asB, Peers: []string{peerB},
				Routes: []routetable.Route{routeOf(peerB, asB, 20)}}},
		},
		{
			name:  "a prefix no peer carries is removed",
			steps: []step{announce(peerA, asA, 10), withdraw(peerA)},
			want:  nil,
		},
		{
			name:  "a withdrawal from a peer that does not carry it changes nothing",
			steps: []step{announce(peerA, asA, 10), withdraw(peerB)},
			want: []routetable.Origin{{ASN: asA, Peers: []string{peerA},
				Routes: []routetable.Route{routeOf(peerA, asA, 10)}}},
		},
		{
			name: "a peer going down leaves the prefix",
			steps: []step{announce(peerA, asA, 10), announce(peerB, asA, 20),
				peerDown(peerB)},
			want: []routetable.Origin{{ASN: asA, Peers: []string{peerA},
				Routes: []routetable.Route{routeOf(peerA, asA, 10)}}},
		},
		{
			name: "a peer going down after it withdrew changes nothing",
			steps: []step{announce(peerA, asA, 10), announce(peerB, asA, 20),
				withdraw(peerB), peerDown(peerB)},
			want: []routetable.Origin{{ASN: asA, Peers: []string{peerA},
				Routes: []routetable.Route{routeOf(peerA, asA, 10)}}},
		},
	})
}

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

	table := routetable.New(routetable.Params{})
	prefixes := []netip.Prefix{
		netip.MustParsePrefix(prefixV4),
		netip.MustParsePrefix("10.2.0.0/16"),
		netip.MustParsePrefix(prefixV6),
	}

	// peerA carries every prefix with origin asA, peerB the first one with
	// origin asB.
	for _, prefix := range prefixes {
		table.Announce(prefix, asA, routeOf(peerA, asA, 10))
	}

	table.Announce(prefixes[0], asB, routeOf(peerB, asB, 20))
	table.PeerDown(peerA)

	want := []routetable.Origin{{ASN: asB, Peers: []string{peerB},
		Routes: []routetable.Route{routeOf(peerB, asB, 20)}}}

	got := table.Origins(prefixes[0])
	if !reflect.DeepEqual(got, want) {
		t.Errorf("Origins(%s) = %+v, want %+v", prefixes[0], got, want)
	}

	for _, prefix := range prefixes[1:] {
		got := table.Origins(prefix)
		if got != nil {
			t.Errorf("Origins(%s) = %+v after the only peer went down, want nil",
				prefix, got)
		}
	}

	counts := table.Counts()
	if counts.IPv4Prefixes != 1 || counts.IPv6Prefixes != 0 ||
		counts.IPv4Routes != 1 || counts.IPv6Routes != 0 {
		t.Errorf("Counts() = %+v, want one IPv4 prefix and route", counts)
	}

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

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

	prefix := netip.MustParsePrefix(prefixV4)
	table := routetable.New(routetable.Params{})

	steps := []struct {
		apply step
		want  int
	}{
		{announce(peerA, asA, 10), 0},
		{announce(peerB, asA, 20), 0},
		{announce(peerC, asA, 30), 0},
		{withdraw(peerC), 0},
		{withdraw(peerB), 1},
		// An update from the last peer is kept.
		{announce(peerA, asA, 40), 0},
		{withdraw(peerA), 0},
		// Once more without kept routes, then without peers: the origin
		// is removed.
		{announce(peerA, asA, 10), 0},
		{announce(peerB, asA, 20), 0},
		{announce(peerC, asA, 30), 0},
		{withdraw(peerB), 0},
		{withdraw(peerC), 1},
		{peerDown(peerA), 0},
	}

	for i, s := range steps {
		s.apply(table, prefix)

		got := table.Counts().OriginsWithoutKeptRoute
		if got != s.want {
			t.Errorf("after step %d, OriginsWithoutKeptRoute = %d, want %d",
				i, got, s.want)
		}
	}
}

// change is a routetable.Change with its peer set as a list of peer numbers.
type change struct {
	Prefix    string
	OriginASN int
	Removed   bool
	Peers     []uint
	Routes    []routetable.Route
}

// peerNumbers returns the numbers in set, or nil for a nil set.
func peerNumbers(set *bitset.BitSet) []uint {
	if set == nil {
		return nil
	}

	return slices.Collect(set.EachSet())
}

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

	// peerA was numbered 5 before; the next new peer is 6.
	table := routetable.New(routetable.Params{
		PeerNumbers: map[string]uint{peerA: 5},
	})
	v4 := netip.MustParsePrefix(prefixV4)
	v6 := netip.MustParsePrefix(prefixV6)

	table.Announce(v6, asA, routeOf(peerA, asA, 10))
	table.Announce(v4, asA, routeOf(peerA, asA, 10))
	table.Announce(v4, asA, routeOf(peerB, asA, 20))

	checkChanges(t, table, []change{
		{Prefix: prefixV4, OriginASN: asA, Peers: []uint{5, 6},
			Routes: []routetable.Route{
				routeOf(peerB, asA, 20), routeOf(peerA, asA, 10),
			}},
		{Prefix: prefixV6, OriginASN: asA, Peers: []uint{5},
			Routes: []routetable.Route{routeOf(peerA, asA, 10)}},
	}, map[string]uint{peerB: 6})

	// TakeChanges cleared the record.
	checkChanges(t, table, []change{}, map[string]uint{})

	// peerB moves to asB; peerA withdraws, which removes the pair of v4 and
	// asA.
	table.Announce(v4, asB, routeOf(peerB, asB, 30))
	table.Withdraw(v4, peerA)

	checkChanges(t, table, []change{
		{Prefix: prefixV4, OriginASN: asA, Removed: true},
		{Prefix: prefixV4, OriginASN: asB, Peers: []uint{6},
			Routes: []routetable.Route{routeOf(peerB, asB, 30)}},
	}, map[string]uint{})
}

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

	table := routetable.New(routetable.Params{})
	v4 := netip.MustParsePrefix(prefixV4)
	v6 := netip.MustParsePrefix(prefixV6)

	// peerA, peerB and peerC are numbered 0, 1 and 2. The routes of peerC and
	// peerB to v4 are kept; peerA keeps only its bit.
	announceAll(table, prefixV4, asA, 10, peerA, peerB, peerC)
	table.Announce(v6, asA, routeOf(peerB, asA, 10))
	table.TakeChanges()

	// peerA announces again, with a time older than both kept routes.
	table.Announce(v4, asA, routeOf(peerA, asA, 15))
	checkChanges(t, table, []change{}, map[string]uint{})

	// peerC announces its kept route again.
	table.Announce(v4, asA, routeOf(peerC, asA, 30))
	checkChanges(t, table, []change{}, map[string]uint{})

	// peerC announces the same path with a newer time, so its kept route
	// changes.
	table.Announce(v4, asA, routeOf(peerC, asA, 40))
	checkChanges(t, table, []change{
		{Prefix: prefixV4, OriginASN: asA, Peers: []uint{0, 1, 2},
			Routes: []routetable.Route{
				routeOf(peerC, asA, 40), routeOf(peerB, asA, 20),
			}},
	}, map[string]uint{})

	// peerA withdraws v6, which it does not carry.
	table.Withdraw(v6, peerA)
	checkChanges(t, table, []change{}, map[string]uint{})

	// peerA going down changes v4, the one prefix it carries.
	table.PeerDown(peerA)
	checkChanges(t, table, []change{
		{Prefix: prefixV4, OriginASN: asA, Peers: []uint{1, 2},
			Routes: []routetable.Route{
				routeOf(peerC, asA, 40), routeOf(peerB, asA, 20),
			}},
	}, map[string]uint{})

	// peerA, which carries nothing now, goes down again.
	table.PeerDown(peerA)
	checkChanges(t, table, []change{}, map[string]uint{})
}

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

	const otherV4 = "10.9.0.0/16"

	v4 := netip.MustParsePrefix(prefixV4)
	v6 := netip.MustParsePrefix(prefixV6)
	other := netip.MustParsePrefix(otherV4)

	// The routes of peerC and peerB to v4 are kept; peerA keeps only its bit.
	// The two routes to v6 have the same time. The routes of other were
	// withdrawn while peerA still carries it.
	source := routetable.New(routetable.Params{})
	announceAll(source, prefixV4, asA, 10, peerA, peerB, peerC)
	announceAll(source, prefixV4, asB, 40, peerD)
	source.Announce(v6, asA, routeOf(peerA, asA, 50))
	source.Announce(v6, asA, routeOf(peerC, asA, 50))
	announceAll(source, otherV4, asB, 10, peerA, peerB, peerC)
	source.Withdraw(other, peerB)
	source.Withdraw(other, peerC)

	changes, peers := source.TakeChanges()
	loaded := routetable.New(routetable.Params{PeerNumbers: peers})

	// Loading a pair twice changes nothing.
	for range 2 {
		for _, change := range changes {
			loaded.Load(change)
		}
	}

	for _, prefix := range []netip.Prefix{v4, v6, other} {
		got, want := loaded.Origins(prefix), source.Origins(prefix)
		if !reflect.DeepEqual(got, want) {
			t.Errorf("Origins(%s) = %+v after the load, want %+v", prefix, got, want)
		}
	}

	if !reflect.DeepEqual(loaded.Counts(), source.Counts()) {
		t.Errorf("Counts() = %+v after the load, want %+v",
			loaded.Counts(), source.Counts())
	}

	checkChanges(t, loaded, []change{}, map[string]uint{})
}

// TestConcurrentUse has goroutines update one table while others read it. The
// race detector, which make test runs, fails it when a method uses the table
// without its lock.
func TestConcurrentUse(t *testing.T) {
	t.Parallel()

	const (
		// rounds is how many times each goroutine calls its method.
		rounds     = 1000
		ipVersion  = 4
		sampleSize = 10
	)

	table := routetable.New(routetable.Params{})
	prefix := netip.MustParsePrefix(prefixV4)

	uses := []func(round int){
		func(round int) { table.Announce(prefix, asA, routeOf(peerA, asA, round)) },
		func(round int) { table.Announce(prefix, asB, routeOf(peerB, asB, round)) },
		func(int) { table.Withdraw(prefix, peerA) },
		func(int) { table.PeerDown(peerB) },
		func(int) { table.TakeChanges() },
		func(int) { table.Lookup(prefix.Addr()) },
		func(int) { table.Origins(prefix) },
		func(int) { table.OriginatedPrefixes(asA) },
		func(int) { table.SamplePrefixes(ipVersion, prefix.Bits(), sampleSize) },
		func(int) { table.Counts() },
	}

	var wg sync.WaitGroup
	for _, use := range uses {
		wg.Go(func() {
			for round := range rounds {
				use(round)
			}
		})
	}

	wg.Wait()
}

// checkChanges checks that TakeChanges returns want and wantPeers.
func checkChanges(
	t *testing.T, table *routetable.Table, want []change,
	wantPeers map[string]uint,
) {
	t.Helper()

	changes, newPeers := table.TakeChanges()

	got := make([]change, 0, len(changes))
	for _, c := range changes {
		got = append(got, change{
			Prefix: c.Prefix.String(), OriginASN: c.OriginASN, Removed: c.Removed,
			Peers: peerNumbers(c.Peers), Routes: c.Routes,
		})
	}

	if !reflect.DeepEqual(got, want) {
		t.Errorf("TakeChanges() changes = %+v, want %+v", got, want)
	}

	if !reflect.DeepEqual(newPeers, wantPeers) {
		t.Errorf("TakeChanges() new peers = %v, want %v", newPeers, wantPeers)
	}
}