// Copyright (c) Tailscale Inc & AUTHORS
// SPDX-License-Identifier: BSD-3-Clause

package connstats

import (
	"context"
	"encoding/binary"
	"fmt"
	"math/rand"
	"net/netip"
	"runtime"
	"sync"
	"testing"
	"time"

	qt "github.com/frankban/quicktest"
	"tailscale.com/cmd/testwrapper/flakytest"
	"tailscale.com/types/ipproto"
	"tailscale.com/types/netlogtype"
)

func testPacketV4(proto ipproto.Proto, srcAddr, dstAddr [4]byte, srcPort, dstPort, size uint16) (out []byte) {
	var ipHdr [20]byte
	ipHdr[0] = 4<<4 | 5
	binary.BigEndian.PutUint16(ipHdr[2:], size)
	ipHdr[9] = byte(proto)
	*(*[4]byte)(ipHdr[12:]) = srcAddr
	*(*[4]byte)(ipHdr[16:]) = dstAddr
	out = append(out, ipHdr[:]...)
	switch proto {
	case ipproto.TCP:
		var tcpHdr [20]byte
		binary.BigEndian.PutUint16(tcpHdr[0:], srcPort)
		binary.BigEndian.PutUint16(tcpHdr[2:], dstPort)
		out = append(out, tcpHdr[:]...)
	case ipproto.UDP:
		var udpHdr [8]byte
		binary.BigEndian.PutUint16(udpHdr[0:], srcPort)
		binary.BigEndian.PutUint16(udpHdr[2:], dstPort)
		out = append(out, udpHdr[:]...)
	default:
		panic(fmt.Sprintf("unknown proto: %d", proto))
	}
	return append(out, make([]byte, int(size)-len(out))...)
}

// TestInterval ensures that we receive at least one call to `dump` using only
// maxPeriod.
func TestInterval(t *testing.T) {
	c := qt.New(t)

	const maxPeriod = 10 * time.Millisecond
	const maxConns = 2048

	gotDump := make(chan struct{}, 1)
	stats := NewStatistics(maxPeriod, maxConns, func(_, _ time.Time, _, _ map[netlogtype.Connection]netlogtype.Counts) {
		select {
		case gotDump <- struct{}{}:
		default:
		}
	})
	defer stats.Shutdown(context.Background())

	srcAddr := netip.AddrFrom4([4]byte{192, 168, 0, byte(rand.Intn(16))})
	dstAddr := netip.AddrFrom4([4]byte{192, 168, 0, byte(rand.Intn(16))})
	srcPort := uint16(rand.Intn(16))
	dstPort := uint16(rand.Intn(16))
	size := uint16(64 + rand.Intn(1024))
	p := testPacketV4(ipproto.TCP, srcAddr.As4(), dstAddr.As4(), srcPort, dstPort, size)
	stats.UpdateRxVirtual(p)

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	select {
	case <-ctx.Done():
		c.Fatal("didn't receive dump within context deadline")
	case <-gotDump:
	}
}

func TestConcurrent(t *testing.T) {
	flakytest.Mark(t, "https://github.com/tailscale/tailscale/issues/7030")
	c := qt.New(t)

	const maxPeriod = 10 * time.Millisecond
	const maxConns = 10
	virtualAggregate := make(map[netlogtype.Connection]netlogtype.Counts)
	stats := NewStatistics(maxPeriod, maxConns, func(start, end time.Time, virtual, physical map[netlogtype.Connection]netlogtype.Counts) {
		c.Assert(start.IsZero(), qt.IsFalse)
		c.Assert(end.IsZero(), qt.IsFalse)
		c.Assert(end.Before(start), qt.IsFalse)
		c.Assert(len(virtual) > 0 && len(virtual) <= maxConns, qt.IsTrue)
		c.Assert(len(physical) == 0, qt.IsTrue)
		for conn, cnts := range virtual {
			virtualAggregate[conn] = virtualAggregate[conn].Add(cnts)
		}
	})
	defer stats.Shutdown(context.Background())
	var wants []map[netlogtype.Connection]netlogtype.Counts
	gots := make([]map[netlogtype.Connection]netlogtype.Counts, runtime.NumCPU())
	var group sync.WaitGroup
	for i := range gots {
		group.Add(1)
		go func(i int) {
			defer group.Done()
			gots[i] = make(map[netlogtype.Connection]netlogtype.Counts)
			rn := rand.New(rand.NewSource(time.Now().UnixNano()))
			var p []byte
			var t netlogtype.Connection
			for j := 0; j < 1000; j++ {
				delay := rn.Intn(10000)
				if p == nil || rn.Intn(64) == 0 {
					proto := ipproto.TCP
					if rn.Intn(2) == 0 {
						proto = ipproto.UDP
					}
					srcAddr := netip.AddrFrom4([4]byte{192, 168, 0, byte(rand.Intn(16))})
					dstAddr := netip.AddrFrom4([4]byte{192, 168, 0, byte(rand.Intn(16))})
					srcPort := uint16(rand.Intn(16))
					dstPort := uint16(rand.Intn(16))
					size := uint16(64 + rand.Intn(1024))
					p = testPacketV4(proto, srcAddr.As4(), dstAddr.As4(), srcPort, dstPort, size)
					t = netlogtype.Connection{Proto: proto, Src: netip.AddrPortFrom(srcAddr, srcPort), Dst: netip.AddrPortFrom(dstAddr, dstPort)}
				}
				t2 := t
				receive := rn.Intn(2) == 0
				if receive {
					t2.Src, t2.Dst = t2.Dst, t2.Src
				}

				cnts := gots[i][t2]
				if receive {
					stats.UpdateRxVirtual(p)
					cnts.RxPackets++
					cnts.RxBytes += uint64(len(p))
				} else {
					cnts.TxPackets++
					cnts.TxBytes += uint64(len(p))
					stats.UpdateTxVirtual(p)
				}
				gots[i][t2] = cnts
				time.Sleep(time.Duration(rn.Intn(1 + delay)))
			}
		}(i)
	}
	group.Wait()
	c.Assert(stats.Shutdown(context.Background()), qt.IsNil)
	wants = append(wants, virtualAggregate)

	got := make(map[netlogtype.Connection]netlogtype.Counts)
	want := make(map[netlogtype.Connection]netlogtype.Counts)
	mergeMaps(got, gots...)
	mergeMaps(want, wants...)
	c.Assert(got, qt.DeepEquals, want)
}

func mergeMaps(dst map[netlogtype.Connection]netlogtype.Counts, srcs ...map[netlogtype.Connection]netlogtype.Counts) {
	for _, src := range srcs {
		for conn, cnts := range src {
			dst[conn] = dst[conn].Add(cnts)
		}
	}
}

func Benchmark(b *testing.B) {
	// TODO: Test IPv6 packets?
	b.Run("SingleRoutine/SameConn", func(b *testing.B) {
		p := testPacketV4(ipproto.UDP, [4]byte{192, 168, 0, 1}, [4]byte{192, 168, 0, 2}, 123, 456, 789)
		b.ResetTimer()
		b.ReportAllocs()
		for range b.N {
			s := NewStatistics(0, 0, nil)
			for j := 0; j < 1e3; j++ {
				s.UpdateTxVirtual(p)
			}
		}
	})
	b.Run("SingleRoutine/UniqueConns", func(b *testing.B) {
		p := testPacketV4(ipproto.UDP, [4]byte{}, [4]byte{}, 0, 0, 789)
		b.ResetTimer()
		b.ReportAllocs()
		for range b.N {
			s := NewStatistics(0, 0, nil)
			for j := 0; j < 1e3; j++ {
				binary.BigEndian.PutUint32(p[20:], uint32(j)) // unique port combination
				s.UpdateTxVirtual(p)
			}
		}
	})
	b.Run("MultiRoutine/SameConn", func(b *testing.B) {
		p := testPacketV4(ipproto.UDP, [4]byte{192, 168, 0, 1}, [4]byte{192, 168, 0, 2}, 123, 456, 789)
		b.ResetTimer()
		b.ReportAllocs()
		for range b.N {
			s := NewStatistics(0, 0, nil)
			var group sync.WaitGroup
			for j := 0; j < runtime.NumCPU(); j++ {
				group.Add(1)
				go func() {
					defer group.Done()
					for k := 0; k < 1e3; k++ {
						s.UpdateTxVirtual(p)
					}
				}()
			}
			group.Wait()
		}
	})
	b.Run("MultiRoutine/UniqueConns", func(b *testing.B) {
		ps := make([][]byte, runtime.NumCPU())
		for i := range ps {
			ps[i] = testPacketV4(ipproto.UDP, [4]byte{192, 168, 0, 1}, [4]byte{192, 168, 0, 2}, 0, 0, 789)
		}
		b.ResetTimer()
		b.ReportAllocs()
		for range b.N {
			s := NewStatistics(0, 0, nil)
			var group sync.WaitGroup
			for j := 0; j < runtime.NumCPU(); j++ {
				group.Add(1)
				go func(j int) {
					defer group.Done()
					p := ps[j]
					j *= 1e3
					for k := 0; k < 1e3; k++ {
						binary.BigEndian.PutUint32(p[20:], uint32(j+k)) // unique port combination
						s.UpdateTxVirtual(p)
					}
				}(j)
			}
			group.Wait()
		}
	})
}