// Package routetable holds the whole routing table in memory. For each prefix
// and origin AS, it keeps the peers whose current route to the prefix has that
// origin AS. It also keeps the two newest of those routes in full. It applies
// announcements, withdrawals and peers' sessions going down. It answers the
// route reads that the pages and the API need.
package routetable

import (
	"cmp"
	"maps"
	"net/netip"
	"slices"
	"sync"
	"time"

	"github.com/bits-and-blooms/bitset"
	"github.com/google/btree"
)

const (
	// keptRoutesPerOrigin is how many routes of each prefix and origin AS are
	// kept in full: the newest ones. Every other peer that carries the prefix
	// with that origin AS keeps only its bit in the peer set.
	keptRoutesPerOrigin = 2

	// keptTimesDegree is the degree of the B-tree that holds the times of the
	// kept routes in order: each of its nodes holds up to 63 times.
	keptTimesDegree = 32

	ipv4Bits = 32
	ipv6Bits = 128

	// The address lookup ignores prefixes shorter than these mask lengths. A
	// default route or a route this short only says where a peer sends traffic
	// it has no better route for; it does not say who holds the address.
	minLookupMaskLengthV4 = 8
	minLookupMaskLengthV6 = 16

	// The IP versions SamplePrefixes takes.
	ipVersion4 = 4
	ipVersion6 = 6
)

// Route is one peer's route to one prefix. The next hop is the peer itself, so
// it is not stored.
type Route struct {
	// Peer is the peer's address.
	Peer   string
	ASPath []int
	// Time is when the peer announced the route.
	Time time.Time
}

// Origin is one origin AS of a prefix: the peers whose current route to the
// prefix has that origin AS, and the newest of those routes in full.
type Origin struct {
	// ASN is the origin AS. 0 means no known origin: the path ends in an AS
	// set of several members.
	ASN int
	// Peers are the peers' addresses, in the order of their numbers: the
	// order they were first seen in.
	Peers []string
	// Routes are the kept routes, newest first: at most two, and none when
	// both were withdrawn while other peers still carry the prefix.
	Routes []Route
}

// Change is the current content of one prefix and origin AS pair that changed:
// what a database row for the pair holds.
type Change struct {
	Prefix    netip.Prefix
	OriginASN int
	// Removed says that no peer carries the prefix with this origin AS any
	// more. Peers and Routes are then nil.
	Removed bool
	// Peers holds the number of each peer that carries the prefix with this
	// origin AS.
	Peers *bitset.BitSet
	// Routes are the kept routes, newest first.
	Routes []Route
}

// Params configures New.
type Params struct {
	// PeerNumbers holds the number each peer was given before, by address. A
	// peer not in it gets the next number above all of them the first time it
	// announces a route.
	PeerNumbers map[string]uint
}

// Table is the routing table. Its methods are safe for concurrent use.
type Table struct {
	mu sync.RWMutex

	ipv4 family
	ipv6 family
	// byOrigin holds, for each origin AS, the prefixes it originates.
	byOrigin map[int]map[netip.Prefix]*origin

	// peerNumbers and peerAddresses give each peer's number by its address
	// and its address by its number. nextPeerNumber is the number the next
	// new peer gets.
	peerNumbers    map[string]uint
	peerAddresses  map[uint]string
	nextPeerNumber uint
	// peerPrefixes[n] counts the prefixes the peer numbered n carries, for
	// each n below nextPeerNumber. It is a slice, not a map, because Load
	// counts once for every peer of every pair it loads: with a map, reading
	// a full table at startup took several times as long.
	peerPrefixes []int

	// newPeers and changed are what TakeChanges hands over: the peers
	// numbered since it last ran, and the prefix and origin AS pairs changed
	// since it last took them.
	newPeers map[string]uint
	changed  map[pair]struct{}

	// keptTimes counts the kept routes at each time, in time order, so its
	// first and last items are the oldest and newest kept routes.
	keptTimes *btree.BTreeG[keptTime]
	// originsWithoutKeptRoute counts the origins that peers carry but that
	// keep no route.
	originsWithoutKeptRoute int
}

// New returns an empty table that numbers new peers on from
// params.PeerNumbers.
func New(params Params) *Table {
	t := &Table{
		ipv4:          newFamily(ipv4Bits, minLookupMaskLengthV4),
		ipv6:          newFamily(ipv6Bits, minLookupMaskLengthV6),
		byOrigin:      make(map[int]map[netip.Prefix]*origin),
		peerNumbers:   make(map[string]uint, len(params.PeerNumbers)),
		peerAddresses: make(map[uint]string, len(params.PeerNumbers)),
		newPeers:      make(map[string]uint),
		changed:       make(map[pair]struct{}),
		keptTimes: btree.NewG(keptTimesDegree, func(a, b keptTime) bool {
			return a.time.Before(b.time)
		}),
	}

	for address, number := range params.PeerNumbers {
		t.peerNumbers[address] = number
		t.peerAddresses[number] = address
		t.nextPeerNumber = max(t.nextPeerNumber, number+1)
	}

	t.peerPrefixes = make([]int, t.nextPeerNumber)

	return t
}

// Announce applies an announcement: the peer's current route to prefix is now
// r, with origin AS originASN. A peer that carried the prefix with another
// origin AS leaves that origin's peer set, and its route there is dropped. r
// is kept in full when it is one of the two newest routes of its prefix and
// origin AS; otherwise only the peer's bit is.
//
// An announcement whose path is empty or holds AS 0 is a withdrawal, and the
// caller passes it to Withdraw instead.
func (t *Table) Announce(prefix netip.Prefix, originASN int, r Route) {
	t.mu.Lock()
	defer t.mu.Unlock()

	if !prefix.IsValid() {
		return
	}

	prefix = prefix.Masked()
	peer := t.numberPeer(r.Peer)

	current := originWithPeer(t.origins(prefix), peer)
	if current != nil && current.asn != originASN {
		t.leave(prefix, current, peer)
	}

	t.join(prefix, t.findOrAddOrigin(prefix, originASN), route{
		peer: peer,
		path: slices.Clone(r.ASPath),
		time: r.Time,
	})
}

// Withdraw applies a withdrawal: the peer no longer carries prefix. It leaves
// the prefix's peer set, and its kept route, if any, is dropped.
func (t *Table) Withdraw(prefix netip.Prefix, peer string) {
	t.mu.Lock()
	defer t.mu.Unlock()

	number, known := t.peerNumbers[peer]
	if !known || !prefix.IsValid() {
		return
	}

	prefix = prefix.Masked()

	o := originWithPeer(t.origins(prefix), number)
	if o != nil {
		t.leave(prefix, o, number)
	}
}

// PeerDown applies a peer's session going down: the peer leaves every peer
// set, and every route of it is dropped. When the peer carries any prefix, it
// visits every prefix in the table.
func (t *Table) PeerDown(peer string) {
	t.mu.Lock()
	defer t.mu.Unlock()

	number, known := t.peerNumbers[peer]
	if !known || t.peerPrefixes[number] == 0 {
		return
	}

	for _, f := range []*family{&t.ipv4, &t.ipv6} {
		for _, prefixes := range f.byLength {
			for prefix, origins := range prefixes {
				o := originWithPeer(origins, number)
				if o != nil {
					t.leave(prefix, o, number)
				}
			}
		}
	}
}

// Load puts one prefix and origin AS pair that the database holds into the
// table, as startup does with each row it reads. The change's peers join the
// pair's peer set, and its routes are kept as an announcement's would be.
// Unlike an update, Load records no change, so TakeChanges does not hand the
// pair back to be written again. Loading a pair twice changes nothing.
//
// The caller makes sure that every peer change names, in its peer set or in
// its routes, has a number in Params.PeerNumbers.
func (t *Table) Load(change Change) {
	t.mu.Lock()
	defer t.mu.Unlock()

	// A removed pair has no peers, and a pair without peers is not in the
	// table.
	if change.Peers == nil || change.Peers.None() || !change.Prefix.IsValid() {
		return
	}

	prefix := change.Prefix.Masked()
	o := t.findOrAddOrigin(prefix, change.OriginASN)
	wasWithoutKeptRoute := o.withoutKeptRoute()

	// The peers not yet in the peer set join it in one union, which takes
	// less time than setting their bits one at a time.
	joining := change.Peers.Difference(&o.peers)
	o.peers.InPlaceUnion(joining)
	t.familyOf(prefix.Addr()).routes += int(joining.Count())

	for peer := range joining.EachSet() {
		t.peerPrefixes[peer]++
	}

	// Oldest first, as they were announced, so that two routes with the same
	// time are kept in the order they had.
	for _, r := range slices.Backward(change.Routes) {
		peer := t.peerNumbers[r.Peer]

		t.dropKeptRoute(o, peer)
		t.keepIfNewest(o, route{
			peer: peer,
			path: slices.Clone(r.ASPath),
			time: r.Time,
		})
	}

	t.countOriginWithoutKeptRoute(o, wasWithoutKeptRoute)
}

// TakeChanges returns at most maxChanges of the prefix and origin AS pairs that
// changed since they were last taken, in prefix order, each with its current
// content. The other changed pairs wait for a later call. It also returns every
// peer numbered since it last ran, by address, so a pair taken later names only
// peers already returned. It then forgets what it returned.
func (t *Table) TakeChanges(maxChanges int) ([]Change, map[string]uint) {
	t.mu.Lock()
	defer t.mu.Unlock()

	changes := make([]Change, 0, min(maxChanges, len(t.changed)))
	for key := range t.changed {
		if len(changes) == maxChanges {
			break
		}

		changes = append(changes, t.currentContent(key))
		delete(t.changed, key)
	}

	// A map keeps the memory of its largest size, so an emptied one is
	// replaced.
	if len(t.changed) == 0 {
		t.changed = make(map[pair]struct{})
	}

	slices.SortFunc(changes, func(a, b Change) int {
		return cmp.Or(a.Prefix.Compare(b.Prefix),
			cmp.Compare(a.OriginASN, b.OriginASN))
	})

	newPeers := t.newPeers
	t.newPeers = make(map[string]uint)

	return changes, newPeers
}

// PutBackChanges records again what one TakeChanges call returned, for a
// caller that could not write it. The next TakeChanges returns those peers
// again, and those pairs wait to be taken again like any changed pair. Each is
// then taken with its content at that time, so a later change of a pair
// replaces the earlier.
func (t *Table) PutBackChanges(changes []Change, newPeers map[string]uint) {
	t.mu.Lock()
	defer t.mu.Unlock()

	for _, change := range changes {
		t.changed[pair{prefix: change.Prefix, asn: change.OriginASN}] = struct{}{}
	}

	maps.Copy(t.newPeers, newPeers)
}

// ChangeCount returns how many changed prefix and origin AS pairs wait for
// TakeChanges.
func (t *Table) ChangeCount() int {
	t.mu.RLock()
	defer t.mu.RUnlock()

	return len(t.changed)
}

// numberPeer returns the peer's number, and gives it the next one the first
// time the table sees it.
func (t *Table) numberPeer(address string) uint {
	number, known := t.peerNumbers[address]
	if known {
		return number
	}

	number = t.nextPeerNumber
	t.nextPeerNumber++
	t.peerNumbers[address] = number
	t.peerAddresses[number] = address
	t.peerPrefixes = append(t.peerPrefixes, 0)
	t.newPeers[address] = number

	return number
}

// familyOf returns the part of the table for addr's IP version.
func (t *Table) familyOf(addr netip.Addr) *family {
	if addr.Is4() {
		return &t.ipv4
	}

	return &t.ipv6
}

// prefixesOfLength returns the prefixes of the IP version and the mask length
// of prefix, a valid masked prefix.
func (t *Table) prefixesOfLength(prefix netip.Prefix) map[netip.Prefix][]*origin {
	return t.familyOf(prefix.Addr()).byLength[prefix.Bits()]
}

// origins returns the origins of prefix, a valid masked prefix.
func (t *Table) origins(prefix netip.Prefix) []*origin {
	return t.prefixesOfLength(prefix)[prefix]
}

// findOrAddOrigin returns the origin of prefix with AS asn, added with no
// peers when the prefix has none with that AS.
func (t *Table) findOrAddOrigin(prefix netip.Prefix, asn int) *origin {
	prefixes := t.prefixesOfLength(prefix)
	origins := prefixes[prefix]

	i, found := slices.BinarySearchFunc(origins, asn,
		func(o *origin, asn int) int { return cmp.Compare(o.asn, asn) })
	if found {
		return origins[i]
	}

	o := &origin{asn: asn}
	prefixes[prefix] = slices.Insert(origins, i, o)

	if t.byOrigin[asn] == nil {
		t.byOrigin[asn] = make(map[netip.Prefix]*origin)
	}

	t.byOrigin[asn][prefix] = o

	return o
}

// removeOrigin removes o, which no peer carries any more, from prefix, and the
// prefix from the table when it has no origin left.
func (t *Table) removeOrigin(prefix netip.Prefix, o *origin) {
	prefixes := t.prefixesOfLength(prefix)

	origins := slices.DeleteFunc(prefixes[prefix],
		func(other *origin) bool { return other == o })
	if len(origins) == 0 {
		delete(prefixes, prefix)
	} else {
		prefixes[prefix] = origins
	}

	delete(t.byOrigin[o.asn], prefix)

	if len(t.byOrigin[o.asn]) == 0 {
		delete(t.byOrigin, o.asn)
	}
}

// join puts the peer of r in the peer set of o, an origin of prefix, and keeps
// r in full when it is one of the two newest routes of o. The peer's earlier
// route, if o keeps one, is replaced: it is no longer the peer's current route.
// The pair is recorded as changed only when its peer set or kept routes change.
func (t *Table) join(prefix netip.Prefix, o *origin, r route) {
	if slices.ContainsFunc(o.routes, r.equal) {
		// The peer announced its kept route again.
		return
	}

	wasWithoutKeptRoute := o.withoutKeptRoute()

	joined := !o.peers.Test(r.peer)
	if joined {
		o.peers.Set(r.peer)

		t.familyOf(prefix.Addr()).routes++
		t.peerPrefixes[r.peer]++
	}

	dropped := t.dropKeptRoute(o, r.peer)
	kept := t.keepIfNewest(o, r)

	if joined || dropped || kept {
		t.recordChange(prefix, o, wasWithoutKeptRoute)
	}
}

// leave takes peer out of the peer set of o, an origin of prefix, and drops
// its kept route, if any.
func (t *Table) leave(prefix netip.Prefix, o *origin, peer uint) {
	wasWithoutKeptRoute := o.withoutKeptRoute()

	o.peers.Clear(peer)
	t.dropKeptRoute(o, peer)

	t.familyOf(prefix.Addr()).routes--
	t.peerPrefixes[peer]--

	t.recordChange(prefix, o, wasWithoutKeptRoute)
}

// keepIfNewest keeps r in o unless both kept routes of o are newer, and says
// whether it kept r. When o already keeps two routes, the older one is
// dropped, and its peer keeps only its bit.
func (t *Table) keepIfNewest(o *origin, r route) bool {
	// The kept routes stay newest first: r goes before the first one that is
	// not newer than it.
	i := 0
	for i < len(o.routes) && o.routes[i].time.After(r.time) {
		i++
	}

	if i == keptRoutesPerOrigin {
		return false
	}

	if len(o.routes) == keptRoutesPerOrigin {
		t.dropKeptRoute(o, o.routes[keptRoutesPerOrigin-1].peer)
	}

	o.routes = slices.Insert(o.routes, i, r)
	t.countKeptTime(r.time, 1)

	return true
}

// dropKeptRoute drops the peer's route from the kept routes of o, if o keeps
// one, and says whether it did.
func (t *Table) dropKeptRoute(o *origin, peer uint) bool {
	i := slices.IndexFunc(o.routes, func(r route) bool { return r.peer == peer })
	if i < 0 {
		return false
	}

	t.countKeptTime(o.routes[i].time, -1)
	o.routes = slices.Delete(o.routes, i, i+1)

	return true
}

// recordChange records that the pair of prefix and the AS of o changed, and
// brings the count of origins without a kept route up to date.
// wasWithoutKeptRoute says whether o was such an origin before the change. An
// origin that no peer carries any more is removed.
func (t *Table) recordChange(
	prefix netip.Prefix, o *origin, wasWithoutKeptRoute bool,
) {
	t.changed[pair{prefix: prefix, asn: o.asn}] = struct{}{}

	t.countOriginWithoutKeptRoute(o, wasWithoutKeptRoute)

	if o.peers.None() {
		t.removeOrigin(prefix, o)
	}
}

// countOriginWithoutKeptRoute brings the count of origins without a kept route
// up to date after a change to o. wasWithoutKeptRoute says whether o was such
// an origin before the change.
func (t *Table) countOriginWithoutKeptRoute(o *origin, wasWithoutKeptRoute bool) {
	if wasWithoutKeptRoute {
		t.originsWithoutKeptRoute--
	}

	if o.withoutKeptRoute() {
		t.originsWithoutKeptRoute++
	}
}

// countKeptTime adds delta to the number of kept routes with time at.
func (t *Table) countKeptTime(at time.Time, delta int) {
	count, _ := t.keptTimes.Get(keptTime{time: at})
	count.time = at
	count.routes += delta

	if count.routes == 0 {
		t.keptTimes.Delete(count)

		return
	}

	t.keptTimes.ReplaceOrInsert(count)
}

// currentContent returns what the table holds for the pair now.
func (t *Table) currentContent(key pair) Change {
	o := t.byOrigin[key.asn][key.prefix]
	if o == nil {
		return Change{Prefix: key.prefix, OriginASN: key.asn, Removed: true}
	}

	return Change{
		Prefix:    key.prefix,
		OriginASN: key.asn,
		Peers:     o.peers.Clone(),
		Routes:    t.exportRoutes(o.routes),
	}
}

// exportRoutes returns copies of routes that name each peer by its address, or
// nil when there are none.
func (t *Table) exportRoutes(routes []route) []Route {
	if len(routes) == 0 {
		return nil
	}

	exported := make([]Route, 0, len(routes))
	for _, r := range routes {
		exported = append(exported, Route{
			Peer:   t.peerAddresses[r.peer],
			ASPath: slices.Clone(r.path),
			Time:   r.time,
		})
	}

	return exported
}

// family is the part of the table for one IP version.
type family struct {
	// byLength[n] holds the prefixes of mask length n that some peer carries,
	// each with its origins in AS number order.
	byLength []map[netip.Prefix][]*origin
	// minLookupLength is the shortest mask length the address lookup uses.
	minLookupLength int
	// routes counts the peer and prefix pairs: each peer that carries a
	// prefix counts once for it.
	routes int
}

// newFamily returns an empty family for addresses of bits bits.
func newFamily(bits, minLookupLength int) family {
	f := family{
		byLength:        make([]map[netip.Prefix][]*origin, bits+1),
		minLookupLength: minLookupLength,
	}

	for n := range f.byLength {
		f.byLength[n] = make(map[netip.Prefix][]*origin)
	}

	return f
}

// prefixesByLength returns how many prefixes of each mask length f holds.
func (f *family) prefixesByLength() []int {
	counts := make([]int, len(f.byLength))
	for n, prefixes := range f.byLength {
		counts[n] = len(prefixes)
	}

	return counts
}

// origin is what the table holds for one prefix and origin AS.
type origin struct {
	asn int
	// peers holds the number of each peer whose current route to the prefix
	// has this origin AS.
	peers bitset.BitSet
	// routes are the newest of those routes, newest first, at most
	// keptRoutesPerOrigin.
	routes []route
}

// withoutKeptRoute says whether peers carry o but o keeps no route.
func (o *origin) withoutKeptRoute() bool {
	return o.peers.Any() && len(o.routes) == 0
}

// newestTime returns the time of the newest kept route of o, or nil when o
// keeps none.
func (o *origin) newestTime() *time.Time {
	if len(o.routes) == 0 {
		return nil
	}

	newest := o.routes[0].time

	return &newest
}

// originWithPeer returns the origin whose peer set holds peer, or nil when
// none does. A peer carries a prefix with one origin AS at most.
func originWithPeer(origins []*origin, peer uint) *origin {
	for _, o := range origins {
		if o.peers.Test(peer) {
			return o
		}
	}

	return nil
}

// route is a kept route. peer is the peer's number.
type route struct {
	peer uint
	path []int
	time time.Time
}

// equal says whether r and other are the same peer's route with the same path
// and time.
func (r route) equal(other route) bool {
	return r.peer == other.peer && r.time.Equal(other.time) &&
		slices.Equal(r.path, other.path)
}

// pair names one prefix and origin AS.
type pair struct {
	prefix netip.Prefix
	asn    int
}

// keptTime counts the kept routes with one time.
type keptTime struct {
	time   time.Time
	routes int
}