package routetable

import (
	"net/netip"
	"runtime"
	"testing"
	"time"
)

// TestWalkPrefixesVisitsEveryPrefix walks a table with more prefixes than two
// of the walk's batches, while another goroutine keeps adding and removing
// other prefixes of the same mask length. The walk lets the lock go between
// batches, so those updates get in while it runs. Every prefix that was in
// the table all along must be visited once.
func TestWalkPrefixesVisitsEveryPrefix(t *testing.T) {
	t.Parallel()

	const (
		asn  = 1000
		peer = "192.0.2.1"
		// prefixCount is more prefixes than two of the walk's batches.
		prefixCount = 2*walkBatch + 1
	)

	routeTime := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
	table := New(Params{})
	prefixes := make([]netip.Prefix, 0, prefixCount)
	addr := netip.MustParseAddr("1.0.0.0")

	for range prefixCount {
		prefix := netip.PrefixFrom(addr, addr.BitLen())
		table.Announce(prefix, asn, Route{Peer: peer, Time: routeTime})
		prefixes = append(prefixes, prefix)
		addr = addr.Next()
	}

	stop := make(chan struct{})
	updated := make(chan struct{})

	go func() {
		defer close(updated)

		other := netip.MustParseAddr("2.0.0.0")

		for {
			select {
			case <-stop:
				return
			default:
			}

			prefix := netip.PrefixFrom(other, other.BitLen())
			table.Announce(prefix, asn, Route{Peer: peer, Time: routeTime})
			table.Withdraw(prefix, peer)

			other = other.Next()
		}
	}()

	visits := make(map[netip.Prefix]int)

	table.walkPrefixes(func(prefix netip.Prefix, _ []*origin) {
		visits[prefix]++
	})

	close(stop)
	<-updated

	for _, prefix := range prefixes {
		if visits[prefix] != 1 {
			t.Fatalf("the walk visited %s %d times, want once", prefix,
				visits[prefix])
		}
	}
}

// TestWalkPrefixesLetsUpdatesIn starts an update during the walk's first
// batch. The update must have been made by the time the walk visits the first
// prefix of its second batch, because the walk lets the lock go in between.
func TestWalkPrefixesLetsUpdatesIn(t *testing.T) {
	t.Parallel()

	const (
		asn  = 1000
		peer = "192.0.2.1"
	)

	routeTime := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
	table := New(Params{})
	addr := netip.MustParseAddr("1.0.0.0")

	for range walkBatch + 1 {
		prefix := netip.PrefixFrom(addr, addr.BitLen())
		table.Announce(prefix, asn, Route{Peer: peer, Time: routeTime})

		addr = addr.Next()
	}

	updatePrefix := netip.MustParsePrefix("2.0.0.0/8")
	updateDone := make(chan struct{})
	visited := 0
	madeBeforeSecondBatch := false

	table.walkPrefixes(func(netip.Prefix, []*origin) {
		visited++

		switch visited {
		case 1:
			go func() {
				defer close(updateDone)

				table.Announce(updatePrefix, asn, Route{Peer: peer, Time: routeTime})
			}()

			// Wait until the update waits for the lock. TryRLock fails
			// once a writer waits.
			for table.mu.TryRLock() {
				table.mu.RUnlock()
				runtime.Gosched()
			}
		case walkBatch + 1:
			// The walk holds the read lock, so the table can be read.
			madeBeforeSecondBatch = table.origins(updatePrefix) != nil
		}
	})

	<-updateDone

	if !madeBeforeSecondBatch {
		t.Fatal("an update started during the walk's first batch was not made " +
			"before its second batch")
	}
}