package routewatch

import (
	"context"
	"net/netip"
	"slices"
	"sync"
	"time"

	"git.eeqj.de/sneak/routewatch/internal/database"
	"git.eeqj.de/sneak/routewatch/internal/logger"
	"git.eeqj.de/sneak/routewatch/internal/metrics"
	"git.eeqj.de/sneak/routewatch/internal/ristypes"
	"git.eeqj.de/sneak/routewatch/internal/routetable"
)

const (
	// prefixHandlerQueueSize is the most messages the handler's queue holds.
	// While the routing table's Overview or PeersDown walks every prefix, it
	// holds the table's lock about half of the time. The handler then applies
	// updates more slowly. When the feed pauses for some seconds and then
	// catches up, its backlog can arrive faster than that, and the queue
	// holds the backlog. The streamer drops an update that finds the queue
	// full. A message takes about 1 KiB, so a full queue takes about 100 MiB
	// (README "Memory").
	prefixHandlerQueueSize = 100000

	// prefixBatchSize is the number of prefix updates to batch together
	prefixBatchSize = 25000

	// prefixBatchMaxSize is the most prefix updates the batch holds while the
	// flush goroutine writes the previous one: ten batches, about 15 MiB
	// (README "Memory"). At this size HandleMessage waits until the flush
	// goroutine takes the batch.
	prefixBatchMaxSize = 250000

	// prefixBatchTimeout is the maximum time to wait before flushing a batch
	// DO NOT reduce this timeout - larger batches are more efficient
	prefixBatchTimeout = 1 * time.Second

	// maxTableChangesPerWrite is the most changed prefix and origin AS pairs
	// of the routing table that one flush writes; the rest wait for the next
	// flushes. The change of one pair takes about 430 bytes of memory and 15
	// microseconds to write. So a write holds about 9 MB and takes about a
	// third of the one-second flush interval, which leaves the rest to the
	// batch. The feed peaks at about 5,000 updates a second, each changing at
	// most one pair. So the million pairs that a full-table peer's session
	// going down changes are written in about a minute.
	maxTableChangesPerWrite = 20000

	// maxOriginChangesPerWrite is the most origin changes of the routing
	// table that one flush writes; the rest wait for the next flushes.
	maxOriginChangesPerWrite = 20000

	// maxHeldBackMessages is the most messages the handler holds back from
	// peers whose session went down. A peer's messages are held back from
	// the moment its session goes down until peerDown has applied them. If
	// peerDown is running at that moment, the peer waits for the rest of
	// peerDown's current PeersDown call. The peer also waits while peerDown
	// applies the messages held back from that call's peers, those that
	// arrive meanwhile included. The next PeersDown call takes the peer out.
	// Once that call returns, peerDown applies the messages held back from
	// that call's peers, one peer after another, in the order the call lists
	// them. A held-back message that says the peer's session went down again
	// makes the peer's later messages wait for the next call. A message takes
	// about 1 KiB, so 150,000 take about 150 MiB (README "Memory"). At this
	// number HandleMessage waits until one of them is applied.
	maxHeldBackMessages = 150000
)

// PrefixHandler keeps the routing table in memory current from the feed. It
// applies each update to the table as it handles it. An announcement gives
// the peer a route to the prefix. A withdrawal removes the peer's route, as
// does an announcement whose AS path is empty or holds AS 0 anywhere. A peer
// whose BGP session goes down leaves every prefix.
//
// The table's PeersDown applies sessions going down. It visits every prefix,
// which takes seconds on a full table. So the handler runs PeersDown on a
// separate goroutine and goes on taking messages meanwhile. One PeersDown runs
// at a time. The next one takes every peer whose session went down since the
// previous one began. PeersDown lets the table's lock go about half of the
// time, so the handler applies other peers' messages while PeersDown runs. The
// handler holds back a down peer's later messages until peerDown applies them,
// as the doc of maxHeldBackMessages describes. It applies them in order, one at
// a time, and takes new messages in between. It holds back at most
// maxHeldBackMessages messages.
//
// Each flush writes the table's changes and origin changes to the database.
// It also stores each prefix the updates name with the time it was first seen.
type PrefixHandler struct {
	db      database.Store
	logger  *logger.Logger
	metrics *metrics.Tracker
	table   *routetable.Table

	// batchMaxSize is prefixBatchMaxSize. Only HandleMessage reads it, so a
	// test may lower it before it hands the handler a message.
	batchMaxSize int
	// tableChangesPerWrite is maxTableChangesPerWrite, except in tests that
	// give a smaller one.
	tableChangesPerWrite int
	// originChangesPerWrite is maxOriginChangesPerWrite, except in tests that
	// give a smaller one.
	originChangesPerWrite int

	// Batching. mu guards batch, writing, writingTableChanges and
	// writingOriginChanges: the flush goroutine takes them under mu and
	// writes them after releasing mu, so HandleMessage waits for a database
	// write only while batch is at batchMaxSize.
	mu    sync.Mutex
	batch []seenPrefix
	// writing is how many updates the flush goroutine took and has not
	// finished writing.
	writing int
	// writingTableChanges is how many changes the flush goroutine took from
	// the routing table and has not finished writing.
	writingTableChanges int
	// writingOriginChanges is how many origin changes the flush goroutine
	// took from the routing table and has not finished writing.
	writingOriginChanges int
	// batchTaken is signalled when the flush goroutine takes batch; its L is
	// &mu.
	batchTaken  sync.Cond
	batchFullCh chan struct{} // wakes the flush goroutine when batch is full
	// stopCh hands Stop's context to the flush goroutine. It holds one, so
	// Stop does not wait for a write in progress to hand it over.
	stopCh chan context.Context
	done   chan struct{} // closed when the flush goroutine returns

	// Peers whose session went down. heldBack, under mu, holds each such
	// peer's later messages, in order, until the peer has left every prefix
	// and they are applied. A peer with no such messages yet has an empty
	// entry. heldCount, under mu, is how many messages heldBack holds.
	// heldTaken is signalled each time a held-back message is taken out to
	// be applied; its L is &mu.
	heldBack  map[string][]*ristypes.RISMessage
	heldCount int
	heldTaken sync.Cond
	// downPeers, under mu, lists the peers whose session went down and that
	// the next PeersDown takes. peerDownRunning, under mu, says whether
	// peerDown is running. peerDowns counts the peerDown goroutines that have
	// not returned.
	downPeers       []string
	peerDownRunning bool
	peerDowns       sync.WaitGroup
	// tablePeersDown is the table's PeersDown. Only peerDown calls it, so a
	// test may replace it before it hands the handler a message.
	tablePeersDown func(peers []string)
}

// seenPrefix is what the batch keeps of an update: the prefix it names and its
// time. A flush needs only these to store each new prefix's first-seen time.
type seenPrefix struct {
	prefix    string
	timestamp time.Time
}

// prefixUpdate is one prefix of an update message, as applyToTable applies it
// to the routing table.
type prefixUpdate struct {
	prefix      string
	originASN   int
	peer        string
	messageType string
	timestamp   time.Time
	path        []int
	// endsInSet says the path ends in an AS set of several members, so it
	// gives no origin.
	endsInSet bool
}

// announcesRoute says whether u gives the peer a route to the prefix: an
// announcement whose path gives an origin, or ends in an AS set of several
// members, which gives origin 0 (none). Any other update removes the peer's
// route for the prefix: a withdrawal, or an announcement whose path is empty
// or holds AS 0 anywhere (the decoder empties such a path). Such an
// announcement replaces the peer's route, as any announcement does, with none
// (RFC 7607 treats AS 0 in a path as a withdrawal).
func (u prefixUpdate) announcesRoute() bool {
	return u.messageType == "announcement" && (u.originASN > 0 || u.endsInSet)
}

// NewPrefixHandler creates a new batched prefix tracking handler, which applies
// the updates to table, the routing table read from db at startup.
func NewPrefixHandler(
	db database.Store, logger *logger.Logger, table *routetable.Table,
) *PrefixHandler {
	return newPrefixHandler(db, logger, table, prefixBatchTimeout,
		maxTableChangesPerWrite)
}

// newPrefixHandler is NewPrefixHandler with the flush timer's interval,
// batchTimeout, in place of prefixBatchTimeout, and tableChangesPerWrite in
// place of maxTableChangesPerWrite. A test that the timer must not disturb
// gives an interval longer than the test runs. A test of writes split over
// several flushes gives a small tableChangesPerWrite, so that it exceeds it
// with few changes.
func newPrefixHandler(
	db database.Store, logger *logger.Logger, table *routetable.Table,
	batchTimeout time.Duration, tableChangesPerWrite int,
) *PrefixHandler {
	h := &PrefixHandler{
		db:                    db,
		logger:                logger,
		table:                 table,
		batch:                 make([]seenPrefix, 0, prefixBatchSize),
		batchMaxSize:          prefixBatchMaxSize,
		tableChangesPerWrite:  tableChangesPerWrite,
		originChangesPerWrite: maxOriginChangesPerWrite,
		batchFullCh:           make(chan struct{}, 1),
		stopCh:                make(chan context.Context, 1),
		done:                  make(chan struct{}),
		heldBack:              make(map[string][]*ristypes.RISMessage),
		tablePeersDown:        table.PeersDown,
	}
	h.batchTaken.L = &h.mu
	h.heldTaken.L = &h.mu

	// Start the flush timer goroutine
	go flushLoop(batchTimeout, h.batchFullCh, h.stopCh, h.done, h.flushBatch)

	return h
}

// SetMetricsTracker sets the metrics tracker for recording route updates
func (h *PrefixHandler) SetMetricsTracker(metrics *metrics.Tracker) {
	h.metrics = metrics
}

// WantsMessage returns true if this handler wants to process messages of the
// given type
func (h *PrefixHandler) WantsMessage(messageType string) bool {
	// UPDATE messages change routes; a STATE message can say a peer's
	// session went down, which removes all of its routes.
	return messageType == ristypes.MessageTypeUpdate ||
		messageType == ristypes.MessageTypeState
}

// QueueCapacity returns the desired queue capacity for this handler
func (h *PrefixHandler) QueueCapacity() int {
	// Batching allows us to use a larger queue
	return prefixHandlerQueueSize
}

// HandleMessage processes a message to track prefix information. A message
// from a peer whose session went down is held back until peerDown applies it.
func (h *PrefixHandler) HandleMessage(msg *ristypes.RISMessage) {
	h.mu.Lock()
	defer h.mu.Unlock()

	// While maxHeldBackMessages messages are held back, wait until one of them
	// is taken out to be applied. Meanwhile the handler takes nothing off its
	// queue, as at batchMaxSize.
	for h.heldCount >= maxHeldBackMessages {
		h.heldTaken.Wait()
	}

	held, goingDown := h.heldBack[msg.Peer]
	if goingDown {
		h.heldBack[msg.Peer] = append(held, msg)
		h.heldCount++

		return
	}

	h.applyMessage(msg)
}

// HeldBackCount returns how many messages the handler holds back from peers
// whose session went down. The status page shows it.
func (h *PrefixHandler) HeldBackCount() int {
	h.mu.Lock()
	defer h.mu.Unlock()

	return h.heldCount
}

// Stop waits for every PeersDown the handler started and for the messages held
// back meanwhile to be applied. It gives up this wait once ctx ends. Then Stop
// stops the flush goroutine and waits for it until ctx ends. The flush
// goroutine makes one last flush before it returns. That flush writes at most
// tableChangesPerWrite of the routing table's changes, originChangesPerWrite of
// its origin changes, and the prefixes of the updates the handler holds. Once
// ctx has ended the flush goroutine skips that last flush, and a write still
// running is lost when the process exits. A failed write of the table's
// changes is not retried either. Stop logs how many updates, table changes and
// origin changes are not written. It also logs how many held-back messages are
// not applied. The caller calls Stop once its last HandleMessage has returned.
func (h *PrefixHandler) Stop(ctx context.Context) {
	peerDownsDone := make(chan struct{})

	go func() {
		h.peerDowns.Wait()
		close(peerDownsDone)
	}()

	select {
	case <-peerDownsDone:
	case <-ctx.Done():
	}

	h.stopCh <- ctx

	select {
	case <-h.done:
	case <-ctx.Done():
	}

	h.mu.Lock()
	unwritten := h.writing + len(h.batch)
	tableChanges := h.writingTableChanges + h.table.ChangeCount()
	heldMessages := h.heldCount
	originChanges := h.writingOriginChanges + h.table.OriginChangeCount()
	h.mu.Unlock()

	if unwritten == 0 && tableChanges == 0 && heldMessages == 0 &&
		originChanges == 0 {
		return
	}

	if ctx.Err() != nil {
		h.logger.Error("Stop deadline passed before the prefix handler wrote "+
			"its updates", "updates", unwritten, "table_changes", tableChanges,
			"held_messages", heldMessages, "origin_changes", originChanges)

		return
	}

	h.logger.Error("The prefix handler stopped before it wrote everything",
		"updates", unwritten, "table_changes", tableChanges,
		"held_messages", heldMessages, "origin_changes", originChanges)
}

// applyMessage applies msg. The caller holds mu. A peer's session going down
// has the peer leave every prefix. An update is applied to the routing table,
// and its prefixes join the batch.
func (h *PrefixHandler) applyMessage(msg *ristypes.RISMessage) {
	if msg.Type == ristypes.MessageTypeState {
		if msg.State == ristypes.StateDown {
			h.startPeerDown(msg.Peer)
		}

		return
	}

	// Use the pre-parsed timestamp
	timestamp := msg.ParsedTimestamp

	originASN := msg.Path.Origin()

	// The stored and shown AS path lists the members of an AS set that ends
	// the path after the rest of it.
	path := slices.Concat(msg.Path.Sequence, msg.Path.Set)

	// While the batch is at its maximum size, wait until the flush goroutine
	// takes it. Meanwhile the handler takes nothing off its queue, which
	// fills, and the streamer drops messages for it.
	for len(h.batch) >= h.batchMaxSize {
		h.batchTaken.Wait()
	}

	// Withdrawals are applied to the routing table before announcements, so
	// a message that lists a prefix as both withdrawn and announced keeps the
	// route (RFC 4271 section 4.3). A withdrawal names only the prefix and
	// peer: the message's AS path, and so its origin AS, belongs to the
	// announced prefixes.
	for _, prefix := range msg.Withdrawals {
		update := prefixUpdate{
			prefix:      canonicalPrefix(prefix),
			peer:        msg.Peer,
			messageType: "withdrawal",
			timestamp:   timestamp,
		}
		h.batch = append(h.batch, seenPrefix{update.prefix, update.timestamp})
		h.applyToTable(update)

		// Record withdrawal in metrics
		if h.metrics != nil {
			h.metrics.RecordWithdrawal()
		}
	}

	for _, announcement := range msg.Announcements {
		for _, prefix := range announcement.Prefixes {
			update := prefixUpdate{
				prefix:      canonicalPrefix(prefix),
				originASN:   originASN,
				peer:        msg.Peer,
				messageType: "announcement",
				timestamp:   timestamp,
				path:        path,
				endsInSet:   len(msg.Path.Set) > 0,
			}
			h.batch = append(h.batch, seenPrefix{update.prefix, update.timestamp})
			h.applyToTable(update)

			// Record announcement in metrics
			if h.metrics != nil {
				h.metrics.RecordAnnouncement()
			}
		}
	}

	// A full batch wakes the flush goroutine. If it is writing, the wake-up
	// waits in batchFullCh and the batch keeps growing until the write ends
	// or the batch reaches batchMaxSize.
	if len(h.batch) >= prefixBatchSize {
		select {
		case h.batchFullCh <- struct{}{}:
		default: // already woken
		}
	}
}

// startPeerDown adds peer to downPeers. It starts peerDown on a new goroutine
// unless peerDown is running. The caller holds mu. The peer's messages are
// held back until peerDown has applied them. Messages held back for an earlier
// session of the peer going down, and not yet applied, stay held back.
func (h *PrefixHandler) startPeerDown(peer string) {
	_, goingDown := h.heldBack[peer]
	if !goingDown {
		h.heldBack[peer] = []*ristypes.RISMessage{}
	}

	h.downPeers = append(h.downPeers, peer)

	if h.peerDownRunning {
		return
	}

	h.peerDownRunning = true
	h.peerDowns.Add(1)

	go h.peerDown()
}

// peerDown takes every peer in downPeers and runs the table's PeersDown for
// them. Then it applies the messages held back from each of those peers. It
// repeats this until downPeers is empty.
func (h *PrefixHandler) peerDown() {
	defer h.peerDowns.Done()

	for {
		h.mu.Lock()

		peers := h.downPeers
		h.downPeers = nil

		if len(peers) == 0 {
			h.peerDownRunning = false
			h.mu.Unlock()

			return
		}

		h.mu.Unlock()

		h.tablePeersDown(peers)

		for _, peer := range peers {
			h.applyHeldBack(peer)
		}
	}
}

// applyHeldBack applies the messages held back from peer, in order. It takes
// mu for one message at a time, so HandleMessage takes new messages in
// between. The peer's new messages join the end of the held-back ones. A
// held-back message that says the peer's session went down again adds the
// peer to downPeers. The rest stay held back until the next PeersDown has
// returned.
func (h *PrefixHandler) applyHeldBack(peer string) {
	for {
		h.mu.Lock()

		held := h.heldBack[peer]
		if len(held) == 0 {
			delete(h.heldBack, peer)
			h.mu.Unlock()

			return
		}

		msg := held[0]
		held[0] = nil // so that the message is freed once applied
		h.heldBack[peer] = held[1:]
		h.heldCount--
		h.heldTaken.Broadcast()
		h.applyMessage(msg)
		h.mu.Unlock()

		if msg.Type == ristypes.MessageTypeState &&
			msg.State == ristypes.StateDown {
			return
		}
	}
}

// takeBatch returns the pending updates and leaves an empty batch in their
// place for HandleMessage to fill. The updates count in writing until
// flushBatch has written them.
func (h *PrefixHandler) takeBatch() []seenPrefix {
	h.mu.Lock()
	defer h.mu.Unlock()

	if len(h.batch) == 0 {
		return nil
	}

	batch := h.batch
	h.writing = len(batch)
	h.batch = make([]seenPrefix, 0, prefixBatchSize)
	h.batchTaken.Broadcast()

	return batch
}

// flushBatch writes the routing table's changes and origin changes, then
// takes the pending batch and stores the prefixes it names, without holding
// the mutex during the writes.
func (h *PrefixHandler) flushBatch() {
	h.writeTableChanges()
	h.writeOriginChanges()

	batch := h.takeBatch()
	if len(batch) == 0 {
		return
	}

	startTime := time.Now()

	// Each prefix the batch names, announced or withdrawn, with the earliest
	// timestamp among its updates: the prefix tables store only first_seen,
	// written when the prefix is new and never changed after.
	prefixesToUpdate := make(map[string]time.Time)

	for _, update := range batch {
		seen, exists := prefixesToUpdate[update.prefix]
		if !exists || update.timestamp.Before(seen) {
			prefixesToUpdate[update.prefix] = update.timestamp
		}
	}

	err := h.db.UpdatePrefixesBatch(prefixesToUpdate)
	if err != nil {
		h.logger.Error("Failed to update prefix batch",
			"error", err, "count", len(prefixesToUpdate))
	}

	h.mu.Lock()
	h.writing = 0
	h.mu.Unlock()

	h.logger.Debug("Flushed prefix batch",
		"batch_size", len(batch),
		"unique_prefixes", len(prefixesToUpdate),
		"duration_ms", time.Since(startTime).Milliseconds(),
	)
}

// applyToTable applies update to the routing table in memory: as an
// announcement when it gives the peer a route, otherwise as a withdrawal. A
// prefix that does not parse is logged and left out. Each announcement counts
// in the IPv4 or IPv6 route update rate.
func (h *PrefixHandler) applyToTable(update prefixUpdate) {
	prefix, err := netip.ParsePrefix(update.prefix)
	if err != nil {
		h.logger.Error("Failed to parse prefix",
			"prefix", update.prefix, "error", err)

		return
	}

	if !update.announcesRoute() {
		h.table.Withdraw(prefix, update.peer)

		return
	}

	if h.metrics != nil {
		if prefix.Addr().Is4() {
			h.metrics.RecordIPv4Update()
		} else {
			h.metrics.RecordIPv6Update()
		}
	}

	h.table.Announce(prefix, update.originASN, routetable.Route{
		Peer:   update.peer,
		ASPath: update.path,
		Time:   update.timestamp,
	})
}

// writeTableChanges writes at most tableChangesPerWrite changes of the routing
// table in memory, and every peer the table numbered since the last write.
// When the write fails, they go back to the table, so that a later write writes
// them again, each pair with its content then.
func (h *PrefixHandler) writeTableChanges() {
	h.mu.Lock()
	changes, newPeers := h.table.TakeChanges(h.tableChangesPerWrite)
	h.writingTableChanges = len(changes)
	h.mu.Unlock()

	if len(changes) == 0 && len(newPeers) == 0 {
		return
	}

	err := h.db.WriteRouteChanges(changes, newPeers)

	h.mu.Lock()
	defer h.mu.Unlock()

	h.writingTableChanges = 0

	if err != nil {
		h.logger.Error("Failed to write the routing table's changes; the next "+
			"flush retries", "changes", len(changes), "new_peers", len(newPeers),
			"error", err)

		h.table.PutBackChanges(changes, newPeers)
	}
}

// writeOriginChanges writes at most originChangesPerWrite of the routing
// table's origin changes. When the write fails, they go back to the table, so
// that a later write writes them again. A dropped change is logged as a
// warning.
func (h *PrefixHandler) writeOriginChanges() {
	h.mu.Lock()
	changes, dropped := h.table.TakeOriginChanges(h.originChangesPerWrite)
	h.writingOriginChanges = len(changes)
	h.mu.Unlock()

	if dropped > 0 {
		h.logger.Warn("Origin changes dropped", "dropped", dropped)
	}

	if len(changes) == 0 {
		return
	}

	err := h.db.WriteOriginChanges(changes)

	h.mu.Lock()
	defer h.mu.Unlock()

	h.writingOriginChanges = 0

	if err != nil {
		h.logger.Error("Failed to write the origin changes; the next flush "+
			"retries", "changes", len(changes), "error", err)

		h.table.PutBackOriginChanges(changes)
	}
}

// canonicalPrefix returns prefix in the text form net/netip prints, the form
// the prefix tables store and the IP lookup reads them in. The feed sends IPv6
// announcements compressed ("2001:db8::/32") but IPv6 withdrawals uncompressed
// ("2001:db8:0:0:0:0:0:0/32"). A prefix that does not parse is returned
// unchanged, and applyToTable logs it.
func canonicalPrefix(prefix string) string {
	p, err := netip.ParsePrefix(prefix)
	if err != nil {
		return prefix
	}

	return p.Masked().String()
}