package database

import (
	"errors"
	"fmt"
	"math/rand/v2"
	"net/netip"
	"os"
	"reflect"
	"slices"
	"testing"
	"time"

	"github.com/bits-and-blooms/bitset"

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

// Peers, origin ASes and prefixes of the routes table tests.
const (
	peerA = "192.0.2.1"
	peerB = "192.0.2.2"
	peerC = "2001:db8::3"
	peerD = "2001:db8::4"
	peerE = "2001:db8::5"

	asA = 64500
	asB = 64501
	// transitAS is the first AS of every test path, the peer's own.
	transitAS = 64496

	tablePrefixV4      = "10.1.2.0/24"
	tablePrefixShortV4 = "10.9.0.0/16"
	tablePrefixV6      = "2001:db8:1::/48"
)

// tableTime returns the time second seconds after the start of every routes
// table test.
func tableTime(second int) time.Time {
	return time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC).
		Add(time.Duration(second) * time.Second)
}

// announceTo has peer announce prefix with origin AS asn second seconds after
// the start.
func announceTo(
	table *routetable.Table, prefix string, asn int, peer string, second int,
) {
	table.Announce(netip.MustParsePrefix(prefix), asn, routetable.Route{
		Peer: peer, ASPath: []int{transitAS, asn}, Time: tableTime(second),
	})
}

// newRoutesTestDatabase returns a new database in a temporary directory.
func newRoutesTestDatabase(t *testing.T) *Database {
	t.Helper()

	db, err := New(&config.Config{StateDir: t.TempDir()}, logger.New())
	if err != nil {
		t.Fatalf("failed to create database: %v", err)
	}

	t.Cleanup(func() { _ = db.Close() })

	return db
}

// writeTableChanges writes what table changed since the last call.
func writeTableChanges(t *testing.T, db *Database, table *routetable.Table) {
	t.Helper()

	changes, newPeers := table.TakeChanges()

	err := db.WriteRouteChanges(changes, newPeers)
	if err != nil {
		t.Fatalf("WriteRouteChanges: %v", err)
	}
}

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

	db := newRoutesTestDatabase(t)
	source := routetable.New(routetable.Params{})

	// Two origins of one prefix: the routes of peerC and peerB are kept for
	// asA, peerA keeps only its bit.
	announceTo(source, tablePrefixV4, asA, peerA, 10)
	announceTo(source, tablePrefixV4, asA, peerB, 20)
	announceTo(source, tablePrefixV4, asA, peerC, 30)
	announceTo(source, tablePrefixV4, asB, peerD, 40)
	// Two routes with the same time.
	announceTo(source, tablePrefixV6, asA, peerA, 50)
	announceTo(source, tablePrefixV6, asA, peerC, 50)
	// An origin whose kept routes were withdrawn while peerA carries it.
	announceTo(source, tablePrefixShortV4, asB, peerA, 10)
	announceTo(source, tablePrefixShortV4, asB, peerB, 20)
	announceTo(source, tablePrefixShortV4, asB, peerC, 30)
	source.Withdraw(netip.MustParsePrefix(tablePrefixShortV4), peerB)
	source.Withdraw(netip.MustParsePrefix(tablePrefixShortV4), peerC)

	writeTableChanges(t, db, source)
	checkReadTable(t, db, source)

	// A peer moves to another origin, one goes down and a new one comes.
	announceTo(source, tablePrefixV4, asB, peerA, 60)
	source.PeerDown(peerC)
	announceTo(source, tablePrefixV6, asB, peerE, 70)

	writeTableChanges(t, db, source)
	checkReadTable(t, db, source)
}

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

	db := newRoutesTestDatabase(t)
	source := routetable.New(routetable.Params{})

	// peerA and peerB carry tablePrefixV4, peerB and peerC tablePrefixV6.
	announceTo(source, tablePrefixV4, asA, peerA, 10)
	announceTo(source, tablePrefixV4, asA, peerB, 20)
	announceTo(source, tablePrefixV6, asA, peerB, 30)
	announceTo(source, tablePrefixV6, asA, peerC, 40)

	writeTableChanges(t, db, source)
	read := checkReadTable(t, db, source)

	// The same updates to both tables: a loaded peer withdraws a prefix,
	// another one goes down, and a new peer comes.
	for _, table := range []*routetable.Table{source, read} {
		table.Withdraw(netip.MustParsePrefix(tablePrefixV4), peerA)
		table.PeerDown(peerB)
		announceTo(table, tablePrefixShortV4, asB, peerD, 50)
	}

	checkSameReads(t, read, source)
}

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

	db := newRoutesTestDatabase(t)
	table := routetable.New(routetable.Params{})

	// Only the routes of peerB and peerC are kept, so only the peer set names
	// peerA.
	announceTo(table, tablePrefixV4, asA, peerA, 10)
	announceTo(table, tablePrefixV4, asA, peerB, 20)
	announceTo(table, tablePrefixV4, asA, peerC, 30)
	writeTableChanges(t, db, table)

	_, err := db.db.ExecContext(t.Context(),
		"DELETE FROM route_peers WHERE address = ?", peerA)
	if err != nil {
		t.Fatalf("deleting a route peer: %v", err)
	}

	_, err = db.ReadRouteTable(t.Context())
	if !errors.Is(err, errUnknownRoutePeer) {
		t.Fatalf("ReadRouteTable = %v, want %v", err, errUnknownRoutePeer)
	}
}

// checkReadTable checks that the table ReadRouteTable reads answers every read
// as want does, and that reading it recorded no change. It returns that table.
func checkReadTable(
	t *testing.T, db *Database, want *routetable.Table,
) *routetable.Table {
	t.Helper()

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

	checkSameReads(t, got, want)

	changes, newPeers := got.TakeChanges()
	if len(changes) != 0 || len(newPeers) != 0 {
		t.Errorf("the read recorded changes %+v and new peers %v, want none",
			changes, newPeers)
	}

	return got
}

// checkSameReads checks that got answers every read as want does, for the
// prefixes and the origin ASes the tests use.
func checkSameReads(t *testing.T, got, want *routetable.Table) {
	t.Helper()

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

	for _, text := range []string{
		tablePrefixV4, tablePrefixShortV4, tablePrefixV6,
	} {
		prefix := netip.MustParsePrefix(text)

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

		gotSummary, gotFound := got.Lookup(prefix.Addr())
		wantSummary, wantFound := want.Lookup(prefix.Addr())

		if gotFound != wantFound || !reflect.DeepEqual(gotSummary, wantSummary) {
			t.Errorf("Lookup(%s) = %+v, %v; want %+v, %v", prefix.Addr(),
				gotSummary, gotFound, wantSummary, wantFound)
		}

		ipVersion := ipVersionV6
		if prefix.Addr().Is4() {
			ipVersion = ipVersionV4
		}

		gotSample := sortedSample(got, ipVersion, prefix.Bits())
		wantSample := sortedSample(want, ipVersion, prefix.Bits())

		if !reflect.DeepEqual(gotSample, wantSample) {
			t.Errorf("SamplePrefixes(%d, %d) = %+v, want %+v",
				ipVersion, prefix.Bits(), gotSample, wantSample)
		}
	}

	for _, asn := range []int{asA, asB} {
		if !reflect.DeepEqual(got.OriginatedPrefixes(asn),
			want.OriginatedPrefixes(asn)) {
			t.Errorf("OriginatedPrefixes(%d) = %+v, want %+v", asn,
				got.OriginatedPrefixes(asn), want.OriginatedPrefixes(asn))
		}
	}
}

// sortedSample returns every prefix of the IP version and mask length in
// table, in prefix order.
func sortedSample(
	table *routetable.Table, ipVersion, maskLength int,
) []routetable.PrefixSummary {
	// More than the tests have of any mask length.
	const all = 100

	sample := table.SamplePrefixes(ipVersion, maskLength, all)
	slices.SortFunc(sample, func(a, b routetable.PrefixSummary) int {
		return a.Prefix.Compare(b.Prefix)
	})

	return sample
}

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

	db := newRoutesTestDatabase(t)
	table := routetable.New(routetable.Params{})

	announceTo(table, tablePrefixV4, asA, peerA, 10)
	announceTo(table, tablePrefixV4, asB, peerB, 20)
	writeTableChanges(t, db, table)

	// The prefix's only peer with asA withdraws it.
	table.Withdraw(netip.MustParsePrefix(tablePrefixV4), peerA)
	writeTableChanges(t, db, table)

	var origins []int

	rows, err := db.db.QueryContext(t.Context(),
		"SELECT origin_asn FROM routes WHERE prefix = ?", tablePrefixV4)
	if err != nil {
		t.Fatalf("reading the routes: %v", err)
	}
	defer func() { _ = rows.Close() }()

	for rows.Next() {
		var origin int

		err := rows.Scan(&origin)
		if err != nil {
			t.Fatalf("reading a route: %v", err)
		}

		origins = append(origins, origin)
	}

	err = rows.Err()
	if err != nil {
		t.Fatalf("reading the routes: %v", err)
	}

	if !slices.Equal(origins, []int{asB}) {
		t.Errorf("rows of %s have origins %v, want only %d",
			tablePrefixV4, origins, asB)
	}
}

// routePeer is one row of route_peers.
type routePeer struct {
	number  uint
	address string
}

// readRoutePeerRows returns the rows of route_peers in number order.
func readRoutePeerRows(t *testing.T, db *Database) []routePeer {
	t.Helper()

	rows, err := db.db.QueryContext(t.Context(),
		"SELECT number, address FROM route_peers ORDER BY number")
	if err != nil {
		t.Fatalf("reading the route peers: %v", err)
	}
	defer func() { _ = rows.Close() }()

	var peers []routePeer

	for rows.Next() {
		var peer routePeer

		err := rows.Scan(&peer.number, &peer.address)
		if err != nil {
			t.Fatalf("reading a route peer: %v", err)
		}

		peers = append(peers, peer)
	}

	err = rows.Err()
	if err != nil {
		t.Fatalf("reading the route peers: %v", err)
	}

	return peers
}

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

	db := newRoutesTestDatabase(t)
	table := routetable.New(routetable.Params{})

	announceTo(table, tablePrefixV4, asA, peerA, 10)
	announceTo(table, tablePrefixV6, asA, peerC, 20)
	writeTableChanges(t, db, table)

	// peerA is stored already; peerB is new.
	announceTo(table, tablePrefixV4, asA, peerA, 30)
	announceTo(table, tablePrefixV4, asA, peerB, 40)
	writeTableChanges(t, db, table)

	want := []routePeer{{0, peerA}, {1, peerC}, {2, peerB}}

	got := readRoutePeerRows(t, db)
	if !slices.Equal(got, want) {
		t.Errorf("route_peers = %+v, want %+v", got, want)
	}
}

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

	db := newRoutesTestDatabase(t)
	table := routetable.New(routetable.Params{})

	announceTo(table, tablePrefixV4, asA, peerA, 10)
	announceTo(table, tablePrefixV6, asA, peerB, 20)

	// peerB's number is missing, so its kept route cannot be written.
	changes, newPeers := table.TakeChanges()
	delete(newPeers, peerB)

	err := db.WriteRouteChanges(changes, newPeers)
	if !errors.Is(err, errUnknownRoutePeer) {
		t.Fatalf("WriteRouteChanges = %v, want %v", err, errUnknownRoutePeer)
	}

	peers := readRoutePeerRows(t, db)
	if len(peers) != 0 {
		t.Errorf("route_peers = %+v after the failed write, want none", peers)
	}

	var routes int

	err = db.db.QueryRowContext(t.Context(),
		"SELECT COUNT(*) FROM routes").Scan(&routes)
	if err != nil {
		t.Fatalf("counting the routes: %v", err)
	}

	if routes != 0 {
		t.Errorf("routes holds %d rows after the failed write, want none", routes)
	}
}

// Parameters for BenchmarkReadRouteTable.
const (
	// benchTablePairs is how many prefix and origin AS pairs the table holds:
	// one origin AS for each prefix. Every benchTableV6Every-th prefix is an
	// IPv6 /48, the others IPv4 /24s.
	benchTablePairs   = 200000
	benchTableV6Every = 5
	// benchTablePeers is how many peers there are, about as many as RIS has.
	// benchTablePeersPerPair is how many of them, picked at random, carry each
	// pair: RIS's 460 million routes over 1.4 million pairs.
	benchTablePeers        = 1400
	benchTablePeersPerPair = 330
	// benchTableBatch is how many changes each write holds.
	benchTableBatch = 10000
)

// BenchmarkReadRouteTable times the read of a routing table of 200,000 prefix
// and origin AS pairs, each carried by 330 of 1,400 peers and keeping two
// routes. It writes that table first. It reports the bytes of the database
// file per pair. It is run by hand, not by make test, for example in the test
// image:
//
//	go test -run '^$' -bench ReadRouteTable -benchtime 5x ./internal/database/
func BenchmarkReadRouteTable(b *testing.B) {
	db, err := New(&config.Config{StateDir: b.TempDir()}, logger.New())
	if err != nil {
		b.Fatalf("failed to create database: %v", err)
	}
	defer func() { _ = db.Close() }()

	//nolint:gosec // The benchmark needs the same table every run.
	rng := rand.New(rand.NewPCG(1, 2))

	newPeers := make(map[string]uint, benchTablePeers)
	for number := range uint(benchTablePeers) {
		newPeers[benchPeerAddress(number)] = number
	}

	for first := 0; first < benchTablePairs; first += benchTableBatch {
		changes := make([]routetable.Change, 0, benchTableBatch)
		for pair := first; pair < first+benchTableBatch; pair++ {
			changes = append(changes, benchTableChange(rng, pair))
		}

		err := db.WriteRouteChanges(changes, newPeers)
		if err != nil {
			b.Fatalf("WriteRouteChanges: %v", err)
		}

		newPeers = nil
	}

	// The WAL is copied into the database file, whose size is reported.
	err = db.Checkpoint(b.Context())
	if err != nil {
		b.Fatalf("Checkpoint: %v", err)
	}

	file, err := os.Stat(db.path)
	if err != nil {
		b.Fatalf("reading the database file's size: %v", err)
	}

	for b.Loop() {
		table, err := db.ReadRouteTable(b.Context())
		if err != nil {
			b.Fatalf("ReadRouteTable: %v", err)
		}

		counts := table.Counts()
		if counts.IPv4Prefixes+counts.IPv6Prefixes != benchTablePairs {
			b.Fatalf("read %d prefixes, want %d",
				counts.IPv4Prefixes+counts.IPv6Prefixes, benchTablePairs)
		}
	}

	// After the loop: b.Loop drops the metrics reported before it starts.
	b.ReportMetric(float64(file.Size())/benchTablePairs, "file-bytes/pair")
}

// benchPeerAddress returns the address of the peer numbered number.
func benchPeerAddress(number uint) string {
	const addressesPerBlock = 256

	return fmt.Sprintf("198.18.%d.%d",
		number/addressesPerBlock, number%addressesPerBlock)
}

// benchTableChange returns the content of the pair numbered pair: its prefix,
// its origin AS, benchTablePeersPerPair peers picked at random, and the routes
// of two of them, each with an AS path of four ASes.
func benchTableChange(rng *rand.Rand, pair int) routetable.Change {
	const (
		firstOrigin  = 64500
		transitASes  = 50000
		olderSeconds = 10
	)

	text := fmt.Sprintf("%d.%d.%d.0/24", 1+pair>>16, (pair>>8)%256, pair%256)
	if pair%benchTableV6Every == 0 {
		text = fmt.Sprintf("2001:%x:%x::/48", pair>>16, pair%65536)
	}

	prefix := netip.MustParsePrefix(text)
	origin := firstOrigin + pair%transitASes
	peers := bitset.New(benchTablePeers)

	for peers.Count() < benchTablePeersPerPair {
		peers.Set(rng.UintN(benchTablePeers))
	}

	change := routetable.Change{Prefix: prefix, OriginASN: origin, Peers: peers}

	// The routes of the first two peers, newest first.
	for i, number := range slices.Collect(peers.EachSet())[:keptRoutesPerRow] {
		change.Routes = append(change.Routes, routetable.Route{
			Peer: benchPeerAddress(number),
			ASPath: []int{
				firstOrigin + rng.IntN(transitASes),
				firstOrigin + rng.IntN(transitASes),
				firstOrigin + rng.IntN(transitASes),
				origin,
			},
			Time: tableTime(pair - i*olderSeconds),
		})
	}

	return change
}