package routewatch

import (
	"bytes"
	"context"
	"encoding/json"
	"log/slog"
	"slices"
	"sync"
	"testing"
	"time"

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

const (
	// handleMessageTimeout is how long a test waits for HandleMessage to
	// return, or for a write to start, before it fails.
	handleMessageTimeout = 2 * time.Second

	// stillWaitingAfter is how long a test watches a HandleMessage that must
	// wait. One that does not wait returns well within it.
	stillWaitingAfter = 100 * time.Millisecond
)

// blockingStore is a mockStore whose batch writes for the AS, peer and prefix
// handlers, delete of a peer's routes and peering write wait until the test
// calls releaseWrites. writeStarted is closed when the first of them starts. It
// records what each batch write was given.
type blockingStore struct {
	*mockStore

	release      chan struct{}
	releaseOnce  sync.Once
	writeStarted chan struct{}
	startOnce    sync.Once

	mu       sync.Mutex
	prefixes map[string]bool
	asns     map[int]bool
	peers    map[string]bool
}

func newBlockingStore() *blockingStore {
	return &blockingStore{
		mockStore:    newMockStore(),
		release:      make(chan struct{}),
		writeStarted: make(chan struct{}),
		prefixes:     make(map[string]bool),
		asns:         make(map[int]bool),
		peers:        make(map[string]bool),
	}
}

func (s *blockingStore) UpsertLiveRouteBatch(routes []*database.LiveRoute) error {
	s.waitForRelease()
	s.mu.Lock()
	defer s.mu.Unlock()

	for _, route := range routes {
		s.prefixes[route.Prefix] = true
	}

	return nil
}

func (s *blockingStore) GetOrCreateASNBatch(asns map[int]database.SeenTimes) error {
	s.waitForRelease()
	s.mu.Lock()
	defer s.mu.Unlock()

	for asn := range asns {
		s.asns[asn] = true
	}

	return nil
}

func (s *blockingStore) UpdatePeerBatch(peers map[string]database.PeerUpdate) error {
	s.waitForRelease()
	s.mu.Lock()
	defer s.mu.Unlock()

	for peerIP := range peers {
		s.peers[peerIP] = true
	}

	return nil
}

func (s *blockingStore) DeleteLiveRoutesOfPeer(peerIP string) (int, error) {
	s.waitForRelease()

	return s.mockStore.DeleteLiveRoutesOfPeer(peerIP)
}

func (s *blockingStore) RecordPeering(
	asA, asB int, seen database.SeenTimes,
) error {
	s.waitForRelease()

	return s.mockStore.RecordPeering(asA, asB, seen)
}

func (s *blockingStore) waitForRelease() {
	s.startOnce.Do(func() { close(s.writeStarted) })
	<-s.release
}

// releaseWrites lets the held writes, and all later ones, finish. A second call
// does nothing.
func (s *blockingStore) releaseWrites() {
	s.releaseOnce.Do(func() { close(s.release) })
}

// stopLog holds what a handler's Stop logs it did not write.
//
//nolint:tagliatelle // field names are the keys the handlers log
type stopLog struct {
	Msg       string `json:"msg"`
	Updates   int    `json:"updates"`
	PeersDown int    `json:"peers_down"`
	Paths     int    `json:"paths"`
}

// newErrorLogger returns a logger that writes only errors, as JSON lines, to
// logs.
func newErrorLogger(logs *bytes.Buffer) *logger.Logger {
	return &logger.Logger{Logger: slog.New(slog.NewJSONHandler(logs,
		&slog.HandlerOptions{Level: slog.LevelError}))}
}

// lastLogged decodes the last line in logs, which is the one Stop logs: a
// handler logs nothing after it.
func lastLogged(t *testing.T, logs *bytes.Buffer) stopLog {
	t.Helper()

	lines := bytes.Split(bytes.TrimSpace(logs.Bytes()), []byte("\n"))

	var logged stopLog

	err := json.Unmarshal(lines[len(lines)-1], &logged)
	if err != nil {
		t.Fatalf("want the error Stop logs, got %q: %v", logs.String(), err)
	}

	return logged
}

// returnsInTime reports whether f returns within handleMessageTimeout.
func returnsInTime(f func()) bool {
	done := make(chan struct{})

	go func() {
		f()
		close(done)
	}()

	select {
	case <-done:
		return true
	case <-time.After(handleMessageTimeout):
		return false
	}
}

// startBlockedWrite hands h the messages in fill, which fill its batch so that
// the handler writes it to store, and waits until that write starts. store
// holds it until the test calls releaseWrites.
func startBlockedWrite(
	t *testing.T,
	store *blockingStore,
	h interface {
		HandleMessage(msg *ristypes.RISMessage)
	},
	fill []*ristypes.RISMessage,
) {
	t.Helper()

	if !returnsInTime(func() {
		for _, msg := range fill {
			h.HandleMessage(msg)
		}
	}) {
		t.Fatal("HandleMessage did not return after filling the batch: " +
			"it waits for the database write")
	}

	select {
	case <-store.writeStarted:
	case <-time.After(handleMessageTimeout):
		t.Fatal("the handler did not write its full batch")
	}
}

// handleWhileAWriteIsBlocked starts a blocked write with the messages in fill
// and hands h the message extra while store holds the write. HandleMessage must
// return each time. Then it lets the writes finish and stops h, which writes
// what it holds.
func handleWhileAWriteIsBlocked(
	t *testing.T,
	store *blockingStore,
	h interface {
		HandleMessage(msg *ristypes.RISMessage)
		Stop(ctx context.Context)
	},
	fill []*ristypes.RISMessage,
	extra *ristypes.RISMessage,
) {
	t.Helper()

	defer h.Stop(t.Context())
	defer store.releaseWrites()

	startBlockedWrite(t, store, h, fill)

	if !returnsInTime(func() { h.HandleMessage(extra) }) {
		t.Fatal("HandleMessage did not return while a database write was blocked")
	}
}

// handleAtMaxBatchSize starts a blocked write with the messages in fill, then,
// while store holds the write, hands h the messages in toMax, which bring its
// next batch to its maximum size, and the message extra. HandleMessage for
// extra must wait until the test lets the write finish, and return after. Then
// it stops h, which writes what it holds.
func handleAtMaxBatchSize(
	t *testing.T,
	store *blockingStore,
	h interface {
		HandleMessage(msg *ristypes.RISMessage)
		Stop(ctx context.Context)
	},
	fill, toMax []*ristypes.RISMessage,
	extra *ristypes.RISMessage,
) {
	t.Helper()

	defer h.Stop(t.Context())
	defer store.releaseWrites()

	startBlockedWrite(t, store, h, fill)

	for _, msg := range toMax {
		h.HandleMessage(msg)
	}

	returned := make(chan struct{})

	go func() {
		h.HandleMessage(extra)
		close(returned)
	}()

	select {
	case <-returned:
		t.Fatal("HandleMessage returned while the batch was at its maximum " +
			"size and a write was blocked")
	case <-time.After(stillWaitingAfter):
	}

	store.releaseWrites()

	select {
	case <-returned:
	case <-time.After(handleMessageTimeout):
		t.Fatal("HandleMessage did not return after the database write finished")
	}
}

// stopWhileAWriteIsBlocked starts a blocked write with the messages in fill,
// hands h the message extra, and stops h with a deadline that passes while
// store holds the write. Stop must return all the same. Then it lets the write
// finish and waits until done, closed when h's flush goroutine returns. Past
// the deadline that goroutine skips its last write, so extra is not written,
// and Stop logs as not written both extra and what the held write took.
func stopWhileAWriteIsBlocked(
	t *testing.T,
	store *blockingStore,
	h interface {
		HandleMessage(msg *ristypes.RISMessage)
		Stop(ctx context.Context)
	},
	done <-chan struct{},
	fill []*ristypes.RISMessage,
	extra *ristypes.RISMessage,
) {
	t.Helper()

	const stopDeadline = 100 * time.Millisecond

	defer store.releaseWrites()

	startBlockedWrite(t, store, h, fill)
	h.HandleMessage(extra)

	ctx, cancel := context.WithTimeout(t.Context(), stopDeadline)
	defer cancel()

	if !returnsInTime(func() { h.Stop(ctx) }) {
		t.Fatalf("Stop with a deadline of %s did not return while a write "+
			"was blocked", stopDeadline)
	}

	store.releaseWrites()

	select {
	case <-done:
	case <-time.After(handleMessageTimeout):
		t.Fatal("the flush goroutine did not return after the write finished")
	}
}

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

	const extraPrefix = "198.51.100.0/24"

	var logs bytes.Buffer

	store := newBlockingStore()
	h := NewPrefixHandler(store, newErrorLogger(&logs))
	fill := slices.Repeat([]string{testV4Prefix}, prefixBatchSize)
	// The peer's session goes down first, so the write held is the delete of
	// its routes, and Stop must log that peer down too.
	stopWhileAWriteIsBlocked(t, store, h, h.done,
		[]*ristypes.RISMessage{sessionDown(), prefixAnnouncement(fill)},
		prefixAnnouncement([]string{extraPrefix}))

	store.mu.Lock()
	defer store.mu.Unlock()

	if store.prefixes[extraPrefix] {
		t.Errorf("%s was written after the stop deadline", extraPrefix)
	}

	logged := lastLogged(t, &logs)
	if logged.Updates != prefixBatchSize+1 || logged.PeersDown != 1 {
		t.Errorf("Stop logged %d updates and %d peer downs not written, "+
			"want %d and 1", logged.Updates, logged.PeersDown, prefixBatchSize+1)
	}
}

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

	// Each message gives an update for its origin AS and one for each AS
	// in its path: fill asnBatchSize+1, extra 2.
	const wantUpdates = asnBatchSize + 1 + 2

	var logs bytes.Buffer

	store := newBlockingStore()
	h := NewASHandler(store, newErrorLogger(&logs))
	fill := ristypes.ASPath{Sequence: slices.Repeat([]int{testASNA}, asnBatchSize)}
	stopWhileAWriteIsBlocked(t, store, h, h.done,
		[]*ristypes.RISMessage{{Path: fill}},
		&ristypes.RISMessage{Path: ristypes.ASPath{Sequence: []int{testASNB}}})

	store.mu.Lock()
	defer store.mu.Unlock()

	if store.asns[testASNB] {
		t.Errorf("AS%d was written after the stop deadline", testASNB)
	}

	if logged := lastLogged(t, &logs); logged.Updates != wantUpdates {
		t.Errorf("Stop logged %d updates not written, want %d",
			logged.Updates, wantUpdates)
	}
}

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

	const otherPeerIP = "2001:db8:ffff::2"

	var logs bytes.Buffer

	store := newBlockingStore()
	h := NewPeerHandler(store, newErrorLogger(&logs))
	stopWhileAWriteIsBlocked(t, store, h, h.done,
		slices.Repeat([]*ristypes.RISMessage{{Peer: testPeerIP}}, peerBatchSize),
		&ristypes.RISMessage{Peer: otherPeerIP})

	store.mu.Lock()
	defer store.mu.Unlock()

	if store.peers[otherPeerIP] {
		t.Errorf("peer %s was written after the stop deadline", otherPeerIP)
	}

	if logged := lastLogged(t, &logs); logged.Updates != peerBatchSize+1 {
		t.Errorf("Stop logged %d updates not written, want %d",
			logged.Updates, peerBatchSize+1)
	}
}

// TestPeeringHandlerStopCountsTheWriteInProgress hands a running peering
// handler one AS path and stops it with a deadline that passes while store
// holds Stop's last write, of that path's peering. Stop must log that path as
// not written.
func TestPeeringHandlerStopCountsTheWriteInProgress(t *testing.T) {
	t.Parallel()

	const stopDeadline = 100 * time.Millisecond

	var logs bytes.Buffer

	store := newBlockingStore()
	defer store.releaseWrites()

	h := NewPeeringHandler(store, newErrorLogger(&logs))
	handle(h, time.Now().UTC(), testASNA, testASNB)

	ctx, cancel := context.WithTimeout(t.Context(), stopDeadline)
	defer cancel()

	if !returnsInTime(func() { h.Stop(ctx) }) {
		t.Fatalf("Stop with a deadline of %s did not return while a write "+
			"was blocked", stopDeadline)
	}

	store.releaseWrites()

	select {
	case <-h.done:
	case <-time.After(handleMessageTimeout):
		t.Fatal("the processing goroutine did not return after the write " +
			"finished")
	}

	if logged := lastLogged(t, &logs); logged.Paths != 1 {
		t.Errorf("Stop logged %d paths not written, want 1", logged.Paths)
	}
}

// prefixAnnouncement returns a message in which testPeerIP announces prefixes
// with origin AS testASNA.
func prefixAnnouncement(prefixes []string) *ristypes.RISMessage {
	return &ristypes.RISMessage{
		Peer: testPeerIP,
		Path: ristypes.ASPath{Sequence: []int{testASNA}},
		Announcements: []ristypes.RISAnnouncement{
			{NextHop: testPeerIP, Prefixes: prefixes},
		},
	}
}

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

	const extraPrefix = "198.51.100.0/24"

	store := newBlockingStore()
	h := NewPrefixHandler(store, logger.New())
	fill := slices.Repeat([]string{testV4Prefix}, prefixBatchSize)
	handleWhileAWriteIsBlocked(t, store, h,
		[]*ristypes.RISMessage{prefixAnnouncement(fill)},
		prefixAnnouncement([]string{extraPrefix}))

	store.mu.Lock()
	defer store.mu.Unlock()

	if !store.prefixes[extraPrefix] {
		t.Errorf("Stop did not write %s, handled while a write was blocked",
			extraPrefix)
	}
}

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

	store := newBlockingStore()
	h := NewASHandler(store, logger.New())
	fill := ristypes.ASPath{Sequence: slices.Repeat([]int{testASNA}, asnBatchSize)}
	handleWhileAWriteIsBlocked(t, store, h,
		[]*ristypes.RISMessage{{Path: fill}},
		&ristypes.RISMessage{Path: ristypes.ASPath{Sequence: []int{testASNB}}})

	store.mu.Lock()
	defer store.mu.Unlock()

	if !store.asns[testASNB] {
		t.Errorf("Stop did not write AS%d, handled while a write was blocked",
			testASNB)
	}
}

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

	const otherPeerIP = "2001:db8:ffff::2"

	store := newBlockingStore()
	h := NewPeerHandler(store, logger.New())
	handleWhileAWriteIsBlocked(t, store, h,
		slices.Repeat([]*ristypes.RISMessage{{Peer: testPeerIP}}, peerBatchSize),
		&ristypes.RISMessage{Peer: otherPeerIP})

	store.mu.Lock()
	defer store.mu.Unlock()

	if !store.peers[otherPeerIP] {
		t.Errorf("Stop did not write peer %s, handled while a write was blocked",
			otherPeerIP)
	}
}

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

	const extraPrefix = "198.51.100.0/24"

	store := newBlockingStore()
	h := NewPrefixHandler(store, logger.New())
	fill := slices.Repeat([]string{testV4Prefix}, prefixBatchSize)
	toMax := slices.Repeat([]string{testV4Prefix}, prefixBatchMaxSize)
	handleAtMaxBatchSize(t, store, h,
		[]*ristypes.RISMessage{prefixAnnouncement(fill)},
		[]*ristypes.RISMessage{prefixAnnouncement(toMax)},
		prefixAnnouncement([]string{extraPrefix}))

	store.mu.Lock()
	defer store.mu.Unlock()

	if !store.prefixes[extraPrefix] {
		t.Errorf("Stop did not write %s, handled after the batch reached its "+
			"maximum size", extraPrefix)
	}
}

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

	store := newBlockingStore()
	h := NewASHandler(store, logger.New())
	fill := ristypes.ASPath{Sequence: slices.Repeat([]int{testASNA}, asnBatchSize)}
	toMax := ristypes.ASPath{Sequence: slices.Repeat([]int{testASNA}, asnBatchMaxSize)}
	handleAtMaxBatchSize(t, store, h,
		[]*ristypes.RISMessage{{Path: fill}},
		[]*ristypes.RISMessage{{Path: toMax}},
		&ristypes.RISMessage{Path: ristypes.ASPath{Sequence: []int{testASNB}}})

	store.mu.Lock()
	defer store.mu.Unlock()

	if !store.asns[testASNB] {
		t.Errorf("Stop did not write AS%d, handled after the batch reached its "+
			"maximum size", testASNB)
	}
}

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

	const otherPeerIP = "2001:db8:ffff::2"

	store := newBlockingStore()
	h := NewPeerHandler(store, logger.New())
	message := []*ristypes.RISMessage{{Peer: testPeerIP}}
	handleAtMaxBatchSize(t, store, h,
		slices.Repeat(message, peerBatchSize),
		slices.Repeat(message, peerBatchMaxSize),
		&ristypes.RISMessage{Peer: otherPeerIP})

	store.mu.Lock()
	defer store.mu.Unlock()

	if !store.peers[otherPeerIP] {
		t.Errorf("Stop did not write peer %s, handled after the batch reached "+
			"its maximum size", otherPeerIP)
	}
}

// TestHandlersWakeTheFlushGoroutineWhenTheBatchIsFull hands each handler, built
// without its flush goroutine so that nothing takes a wake-up, one update short
// of a full batch, and a second one exactly a full batch. Only the full batch
// leaves a wake-up in batchFullCh.
func TestHandlersWakeTheFlushGoroutineWhenTheBatchIsFull(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name      string
		batchSize int
		// handle hands a new handler n updates and returns its batchFullCh.
		handle func(n int) chan struct{}
	}{
		{
			name:      "prefix",
			batchSize: prefixBatchSize,
			handle: func(n int) chan struct{} {
				h := &PrefixHandler{batchFullCh: make(chan struct{}, 1)}
				h.HandleMessage(prefixAnnouncement(
					slices.Repeat([]string{testV4Prefix}, n)))

				return h.batchFullCh
			},
		},
		{
			name:      "AS",
			batchSize: asnBatchSize,
			handle: func(n int) chan struct{} {
				h := &ASHandler{batchFullCh: make(chan struct{}, 1)}
				// The origin AS and each AS in the path are one update each.
				h.HandleMessage(&ristypes.RISMessage{
					Path: ristypes.ASPath{Sequence: slices.Repeat([]int{testASNA}, n-1)},
				})

				return h.batchFullCh
			},
		},
		{
			name:      "peer",
			batchSize: peerBatchSize,
			handle: func(n int) chan struct{} {
				h := &PeerHandler{batchFullCh: make(chan struct{}, 1)}
				for range n {
					h.HandleMessage(&ristypes.RISMessage{Peer: testPeerIP})
				}

				return h.batchFullCh
			},
		},
	}

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

			got := len(tt.handle(tt.batchSize - 1))
			if got != 0 {
				t.Errorf("one update short of a full batch left %d wake-ups, "+
					"want none", got)
			}

			got = len(tt.handle(tt.batchSize))
			if got != 1 {
				t.Errorf("a full batch left %d wake-ups, want 1", got)
			}
		})
	}
}