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)
}

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

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

	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))
		}
	}

	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)
	}
}

// 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
}