package routewatch

import (
	"bytes"
	"encoding/json"
	"errors"
	"log/slog"
	"net/netip"
	"reflect"
	"slices"
	"strings"
	"testing"
	"time"

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

// More peers and prefixes for the tests of the routing table in memory.
const (
	secondPeerIP = "2001:db8:ffff::2"
	thirdPeerIP  = "2001:db8:ffff::3"
	fourthPeerIP = "192.0.2.4"
	fifthPeerIP  = "192.0.2.5"

	otherV4Prefix = "198.51.100.0/24"
	setV4Prefix   = "203.0.113.0/24"

	// firstManyPrefix is the first prefix announceManyPrefixes announces.
	firstManyPrefix = "10.0.0.0/32"

	// testMaxTableChangesPerWrite is a limit of one write of the routing
	// table's changes that a test exceeds quickly, in place of
	// maxTableChangesPerWrite.
	testMaxTableChangesPerWrite = 10
)

// tableCase is one test of the routing table in memory: the messages handed to
// the prefix handler, and for each prefix the peers that carry it with each
// origin AS after them.
type tableCase struct {
	name     string
	messages []*ristypes.RISMessage
	want     map[string]map[int][]string
}

// TestPrefixHandlerAppliesUpdatesToTheRoutingTable hands the prefix handler
// announcements, withdrawals and a peer's session going down.
func TestPrefixHandlerAppliesUpdatesToTheRoutingTable(t *testing.T) {
	t.Parallel()

	announceAndWithdraw := peerAnnouncement(testPeerIP, testV4Prefix)
	announceAndWithdraw.Withdrawals = []string{testV4Prefix}

	checkTableCases(t, []tableCase{
		{
			name:     "an announcement",
			messages: []*ristypes.RISMessage{peerAnnouncement(testPeerIP, testV4Prefix)},
			want:     map[string]map[int][]string{testV4Prefix: {testASNB: {testPeerIP}}},
		},
		{
			name: "a withdrawal",
			messages: []*ristypes.RISMessage{
				peerAnnouncement(testPeerIP, testV4Prefix),
				peerWithdrawal(testPeerIP, testV4Prefix),
			},
			want: map[string]map[int][]string{testV4Prefix: {}},
		},
		{
			// RFC 4271 section 4.3: the announcement counts, not the withdrawal.
			name:     "a message that announces and withdraws the prefix",
			messages: []*ristypes.RISMessage{announceAndWithdraw},
			want:     map[string]map[int][]string{testV4Prefix: {testASNB: {testPeerIP}}},
		},
		{
			name: "a peer's session going down",
			messages: []*ristypes.RISMessage{
				peerAnnouncement(testPeerIP, testV4Prefix, testV6Prefix),
				peerAnnouncement(secondPeerIP, testV4Prefix),
				sessionDown(),
			},
			want: map[string]map[int][]string{
				testV4Prefix: {testASNB: {secondPeerIP}},
				testV6Prefix: {},
			},
		},
		{
			name: "an announcement after the peer's session went down",
			messages: []*ristypes.RISMessage{
				peerAnnouncement(testPeerIP, testV4Prefix),
				sessionDown(),
				peerAnnouncement(testPeerIP, otherV4Prefix),
			},
			want: map[string]map[int][]string{
				testV4Prefix:  {},
				otherV4Prefix: {testASNB: {testPeerIP}},
			},
		},
	})
}

// TestPrefixHandlerAppliesAPathWithoutAnOriginToTheRoutingTable hands the
// prefix handler announcements whose AS path gives no origin. One whose path
// is empty or holds AS 0 removes the peer's route, as a withdrawal does; one
// whose path ends in an AS set of several members is kept with origin 0.
func TestPrefixHandlerAppliesAPathWithoutAnOriginToTheRoutingTable(t *testing.T) {
	t.Parallel()

	announceWithPath := func(path ristypes.ASPath) *ristypes.RISMessage {
		msg := peerAnnouncement(testPeerIP, testV4Prefix)
		msg.Path = path

		return msg
	}

	checkTableCases(t, []tableCase{
		{
			name: "a path whose last AS is 0",
			messages: []*ristypes.RISMessage{
				peerAnnouncement(testPeerIP, testV4Prefix),
				announceWithPath(ristypes.ASPath{Sequence: []int{testASNA, 0}}),
			},
			want: map[string]map[int][]string{testV4Prefix: {}},
		},
		{
			name: "an empty path",
			messages: []*ristypes.RISMessage{
				peerAnnouncement(testPeerIP, testV4Prefix),
				announceWithPath(ristypes.ASPath{}),
			},
			want: map[string]map[int][]string{testV4Prefix: {}},
		},
		{
			name: "a path with AS 0 before its end",
			messages: []*ristypes.RISMessage{
				peerAnnouncement(testPeerIP, testV4Prefix),
				announceWithPath(decodePath(t, testASNA, 0, testASNB)),
			},
			want: map[string]map[int][]string{testV4Prefix: {}},
		},
		{
			name: "a path that ends in an AS set",
			messages: []*ristypes.RISMessage{
				announceWithPath(decodePath(t, testASNA, []int{testASNC, testASND})),
			},
			want: map[string]map[int][]string{testV4Prefix: {0: {testPeerIP}}},
		},
	})
}

// TestPrefixHandlerKeepsThePathOfAPathEndingInASet announces a prefix with AS
// paths that end in an AS set. Any member of a set of several may be the
// origin, so the routing table keeps the route with origin 0 (none); a set of
// one member counts as that AS. The kept AS path lists the set's members after
// the rest of the path.
func TestPrefixHandlerKeepsThePathOfAPathEndingInASet(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name       string
		path       []any
		wantOrigin int
		wantPath   []int
	}{
		{
			name:     "a set of several members",
			path:     []any{testASNA, testASNB, []int{testASNC, testASND}},
			wantPath: []int{testASNA, testASNB, testASNC, testASND},
		},
		{
			name:       "a set of one member",
			path:       []any{testASNA, testASNB, []int{testASNC}},
			wantOrigin: testASNC,
			wantPath:   []int{testASNA, testASNB, testASNC},
		},
	}

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

			msg := peerAnnouncement(testPeerIP, testV4Prefix)
			msg.Path = decodePath(t, tt.path...)

			h := prefixHandlerWithoutTimer(newMockStore())
			h.HandleMessage(msg)

			origins := h.table.Origins(netip.MustParsePrefix(testV4Prefix))
			if len(origins) != 1 || origins[0].ASN != tt.wantOrigin ||
				len(origins[0].Routes) != 1 ||
				!slices.Equal(origins[0].Routes[0].ASPath, tt.wantPath) {
				t.Errorf("Origins(%s) = %+v, want one route with origin AS%d "+
					"and AS path %v", testV4Prefix, origins, tt.wantOrigin,
					tt.wantPath)
			}
		})
	}
}

// checkTableCases hands the messages of each case to a new prefix handler and
// checks its routing table right after, with no flush.
func checkTableCases(t *testing.T, tests []tableCase) {
	t.Helper()

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

			h := prefixHandlerWithoutTimer(newMockStore())
			for _, msg := range tt.messages {
				h.HandleMessage(msg)
			}

			for prefix, want := range tt.want {
				got := tableOrigins(h.table, prefix)
				if !reflect.DeepEqual(got, want) {
					t.Errorf("peers of %s by origin AS = %v, want %v",
						prefix, got, want)
				}
			}
		})
	}
}

// tableWritesStore is a real database that hands each write of the routing
// table's changes to written once it is made.
type tableWritesStore struct {
	*database.Database

	written chan []routetable.Change
}

func (s tableWritesStore) WriteRouteChanges(
	changes []routetable.Change, newPeers map[string]uint,
) error {
	err := s.Database.WriteRouteChanges(changes, newPeers)
	s.written <- changes

	return err
}

// TestPrefixHandlerWritesTheTableChangesEachSecond hands a running prefix
// handler an announcement and then its withdrawal, each far short of a full
// batch. The flush timer must write the change of each to the database.
func TestPrefixHandlerWritesTheTableChangesEachSecond(t *testing.T) {
	t.Parallel()

	// The two writes the test waits for, and room for one more.
	const writesBuffered = 3

	store := tableWritesStore{
		Database: newTestDatabase(t),
		written:  make(chan []routetable.Change, writesBuffered),
	}
	h := NewPrefixHandler(store, logger.New(), routetable.New(routetable.Params{}))

	defer h.Stop(t.Context())

	prefix := netip.MustParsePrefix(testV4Prefix)

	for _, step := range []struct {
		msg         *ristypes.RISMessage
		wantRemoved bool
	}{
		{msg: peerAnnouncement(testPeerIP, testV4Prefix)},
		{msg: peerWithdrawal(testPeerIP, testV4Prefix), wantRemoved: true},
	} {
		h.HandleMessage(step.msg)

		select {
		case changes := <-store.written:
			if len(changes) != 1 || changes[0].Prefix != prefix ||
				changes[0].OriginASN != testASNB ||
				changes[0].Removed != step.wantRemoved {
				t.Fatalf("wrote %+v, want the pair of %s and AS%d with "+
					"removed %t", changes, prefix, testASNB, step.wantRemoved)
			}
		case <-time.After(stepTimeout):
			t.Fatalf("the routing table's change was not written within %s",
				stepTimeout)
		}
	}
}

// errTestTableWrite is the error a test store's failing write of the routing
// table's changes returns.
var errTestTableWrite = errors.New("test: the write of the table's changes failed")

// firstTableWriteFailsStore is a real database whose first write of the
// routing table's changes fails, writing nothing. writes counts the writes
// asked for.
type firstTableWriteFailsStore struct {
	*database.Database

	writes int
}

func (s *firstTableWriteFailsStore) WriteRouteChanges(
	changes []routetable.Change, newPeers map[string]uint,
) error {
	s.writes++
	if s.writes == 1 {
		return errTestTableWrite
	}

	return s.Database.WriteRouteChanges(changes, newPeers)
}

// TestPrefixHandlerWritesAFailedTableWriteAgain runs two flushes into a
// database whose first write of the routing table's changes fails. Before the
// first, two new peers announce a prefix each; before the second, one of them
// announces its prefix again with another path to the same origin AS. The
// second write must hold the changes and new peers of the first, with the later
// path, so the database holds the table that is in memory.
func TestPrefixHandlerWritesAFailedTableWriteAgain(t *testing.T) {
	t.Parallel()

	db := newTestDatabase(t)
	store := &firstTableWriteFailsStore{Database: db}
	h := prefixHandlerWithoutTimer(store)

	first := peerAnnouncement(testPeerIP, testV4Prefix)
	h.HandleMessage(first)
	h.HandleMessage(peerAnnouncement(secondPeerIP, otherV4Prefix))
	h.flushBatch()

	laterPath := []int{testASNC, testASNB}
	later := peerAnnouncement(testPeerIP, testV4Prefix)
	later.Path = ristypes.ASPath{Sequence: laterPath}
	later.ParsedTimestamp = first.ParsedTimestamp.Add(secondUpdateAfter)
	h.HandleMessage(later)
	h.flushBatch()

	if store.writes != 2 {
		t.Fatalf("the store was asked for %d writes of the table's changes, "+
			"want 2", store.writes)
	}

	read, err := db.ReadRouteTable(t.Context())
	if err != nil {
		t.Fatalf("ReadRouteTable: %v", err)
	}

	checkSameTable(t, read, h.table, testV4Prefix, otherV4Prefix)

	origins := read.Origins(netip.MustParsePrefix(testV4Prefix))
	if len(origins) != 1 || len(origins[0].Routes) != 1 ||
		!slices.Equal(origins[0].Routes[0].ASPath, laterPath) {
		t.Errorf("Origins(%s) = %+v, want one route with the path %v",
			testV4Prefix, origins, laterPath)
	}
}

// tableWriteSizesStore is a real database that records how many changes and
// how many new peers each write of the routing table's changes holds.
type tableWriteSizesStore struct {
	*database.Database

	changes  []int
	newPeers []int
}

func (s *tableWriteSizesStore) WriteRouteChanges(
	changes []routetable.Change, newPeers map[string]uint,
) error {
	s.changes = append(s.changes, len(changes))
	s.newPeers = append(s.newPeers, len(newPeers))

	return s.Database.WriteRouteChanges(changes, newPeers)
}

// TestPrefixHandlerWritesManyTableChangesOverSeveralFlushes has one new peer
// change one and a half times as many prefix and origin AS pairs as one write
// takes, then flushes three times. The first write must hold as many pairs as
// one write takes, and the peer. The second must hold the other pairs, and the
// third flush must write nothing. The database must then hold the table in
// memory. The handler's limit is small, and its timer does not fire.
func TestPrefixHandlerWritesManyTableChangesOverSeveralFlushes(t *testing.T) {
	t.Parallel()

	const (
		pairs   = testMaxTableChangesPerWrite * 3 / 2
		flushes = 3
	)

	db := newTestDatabase(t)
	store := &tableWriteSizesStore{Database: db}
	table := routetable.New(routetable.Params{})
	announceManyPrefixes(table, pairs)

	h := newPrefixHandler(store, logger.New(), table, longerThanAnyTest,
		testMaxTableChangesPerWrite)

	defer stopInTime(t, h)

	for range flushes {
		h.flushBatch()
	}

	wantChanges := []int{
		testMaxTableChangesPerWrite, pairs - testMaxTableChangesPerWrite,
	}
	wantNewPeers := []int{1, 0}

	if !slices.Equal(store.changes, wantChanges) ||
		!slices.Equal(store.newPeers, wantNewPeers) {
		t.Errorf("the writes held %v pairs and %v new peers, want %v and %v",
			store.changes, store.newPeers, wantChanges, wantNewPeers)
	}

	read, err := db.ReadRouteTable(t.Context())
	if err != nil {
		t.Fatalf("ReadRouteTable: %v", err)
	}

	checkSameTable(t, read, h.table, firstManyPrefix)
}

// TestPrefixHandlerStopLogsTheTableChangesLeftUnwritten stops a prefix handler,
// with no deadline, whose routing table holds more changed pairs than one write
// takes. Its last flush writes one write's worth, so Stop must log the pairs
// that no write took, and must not say that the stop deadline passed. The
// flush timer must not fire, or its writes could leave no pairs for Stop to
// log. The handler's limit of one write is small.
func TestPrefixHandlerStopLogsTheTableChangesLeftUnwritten(t *testing.T) {
	t.Parallel()

	const pairs = testMaxTableChangesPerWrite + 1

	var logs bytes.Buffer

	store := &tableWriteSizesStore{Database: newTestDatabase(t)}
	table := routetable.New(routetable.Params{})
	announceManyPrefixes(table, pairs)

	h := newPrefixHandler(store, newErrorLogger(&logs), table, longerThanAnyTest,
		testMaxTableChangesPerWrite)
	h.Stop(t.Context())

	written := 0
	for _, n := range store.changes {
		written += n
	}

	logged := lastLogged(t, &logs)
	if strings.Contains(logged.Msg, "deadline") {
		t.Errorf("Stop logged %q although its deadline had not passed",
			logged.Msg)
	}

	if logged.TableChanges != pairs-written {
		t.Errorf("Stop logged %d table changes not written, want %d",
			logged.TableChanges, pairs-written)
	}
}

// readLog holds what readRouteTable logs.
//
//nolint:tagliatelle // field names are the keys readRouteTable logs
type readLog struct {
	Msg        string `json:"msg"`
	Pairs      int    `json:"pairs"`
	Peers      int    `json:"peers"`
	DurationMs *int64 `json:"duration_ms"`
}

// TestARestartReadsTheRoutingTableBack runs announcements, a withdrawal and a
// peer's session going down through a prefix handler into a database, then
// reads the table back as startup does. The table read must answer as the one
// in memory, and the log must give its pairs and peers. A new handler on the
// table read then takes later updates and writes them as it stops: the
// database must then hold what the first table holds after the same updates.
func TestARestartReadsTheRoutingTableBack(t *testing.T) {
	t.Parallel()

	prefixes := []string{testV4Prefix, testV6Prefix, otherV4Prefix, setV4Prefix}

	db := newTestDatabase(t)
	h := prefixHandlerWithoutTimer(db)

	viaOtherOrigin := peerAnnouncement(thirdPeerIP, testV4Prefix)
	viaOtherOrigin.Path = ristypes.ASPath{Sequence: []int{testASNA, testASNC}}
	endsInSet := peerAnnouncement(secondPeerIP, setV4Prefix)
	endsInSet.Path = decodePath(t, testASNA, []int{testASNC, testASND})

	for _, msg := range []*ristypes.RISMessage{
		peerAnnouncement(testPeerIP, testV4Prefix, testV6Prefix),
		peerAnnouncement(secondPeerIP, testV4Prefix),
		viaOtherOrigin,
		endsInSet,
		peerAnnouncement(thirdPeerIP, otherV4Prefix),
		peerWithdrawal(thirdPeerIP, otherV4Prefix),
		peerAnnouncement(fourthPeerIP, otherV4Prefix),
		sessionDownOf(fourthPeerIP),
	} {
		h.HandleMessage(msg)
	}

	h.flushBatch()

	var logs bytes.Buffer

	table, err := readRouteTable(db,
		&logger.Logger{Logger: slog.New(slog.NewJSONHandler(&logs, nil))})
	if err != nil {
		t.Fatalf("readRouteTable: %v", err)
	}

	checkSameTable(t, table, h.table, prefixes...)

	// testV4Prefix with two origin ASes, testV6Prefix and setV4Prefix; four
	// peers, the last of which carries nothing.
	const wantPairs, wantPeers = 4, 4

	var logged readLog

	err = json.Unmarshal(bytes.TrimSpace(logs.Bytes()), &logged)
	if err != nil || logged.Pairs != wantPairs || logged.Peers != wantPeers ||
		logged.DurationMs == nil {
		t.Errorf("readRouteTable logged %q, want %d pairs, %d peers and the "+
			"duration", logs.String(), wantPairs, wantPeers)
	}

	restarted := NewPrefixHandler(db, logger.New(), table)

	for _, msg := range []*ristypes.RISMessage{
		peerWithdrawal(testPeerIP, testV6Prefix),
		sessionDownOf(secondPeerIP),
		peerAnnouncement(fourthPeerIP, otherV4Prefix),
		peerAnnouncement(fifthPeerIP, setV4Prefix),
	} {
		// The first handler only applies the updates to its table.
		h.HandleMessage(msg)
		restarted.HandleMessage(msg)
	}

	restarted.Stop(t.Context())

	read, err := db.ReadRouteTable(t.Context())
	if err != nil {
		t.Fatalf("ReadRouteTable after the restart: %v", err)
	}

	checkSameTable(t, read, h.table, prefixes...)
}

// peerWithdrawal returns an UPDATE message in which peer withdraws prefixes.
func peerWithdrawal(peer string, prefixes ...string) *ristypes.RISMessage {
	return &ristypes.RISMessage{
		Type:            ristypes.MessageTypeUpdate,
		Peer:            peer,
		ParsedTimestamp: time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC),
		Withdrawals:     prefixes,
	}
}

// announceManyPrefixes has testPeerIP announce n prefixes in table, with the
// path of peerAnnouncement: the /32 of each address from firstManyPrefix on.
func announceManyPrefixes(table *routetable.Table, n int) {
	addr := netip.MustParsePrefix(firstManyPrefix).Addr()
	route := routetable.Route{
		Peer:   testPeerIP,
		ASPath: []int{testASNA, testASNB},
		Time:   time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC),
	}

	for range n {
		table.Announce(netip.PrefixFrom(addr, addr.BitLen()), testASNB, route)
		addr = addr.Next()
	}
}

// sessionDownOf returns the STATE message that says the BGP session of peer
// went down.
func sessionDownOf(peer string) *ristypes.RISMessage {
	return &ristypes.RISMessage{
		Type:  ristypes.MessageTypeState,
		Peer:  peer,
		State: ristypes.StateDown,
	}
}

// tableOrigins returns, for each origin AS of prefix in table, the peers that
// carry the prefix with it.
func tableOrigins(table *routetable.Table, prefix string) map[int][]string {
	origins := make(map[int][]string)
	for _, o := range table.Origins(netip.MustParsePrefix(prefix)) {
		origins[o.ASN] = o.Peers
	}

	return origins
}

// checkSameTable checks that got answers as want does for each of prefixes, and
// that the two tables have the same counts.
func checkSameTable(
	t *testing.T, got, want *routetable.Table, prefixes ...string,
) {
	t.Helper()

	for _, text := range prefixes {
		prefix := netip.MustParsePrefix(text)

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

	gotCounts, wantCounts := got.Counts(), want.Counts()
	if !reflect.DeepEqual(gotCounts, wantCounts) {
		t.Errorf("Counts() = %+v, want %+v", gotCounts, wantCounts)
	}
}