package engine

import (
	"testing"

	"git.difuse.io/Difuse/Mellaris/analyzer"
	"git.difuse.io/Difuse/Mellaris/ruleset"

	"github.com/bwmarrin/snowflake"
)

type legacyUDPStreamValue struct {
	Tuple udpTupleKey
}

type emptyRuleset struct{}

func (emptyRuleset) Analyzers(ruleset.StreamInfo) []analyzer.Analyzer { return nil }
func (emptyRuleset) Match(ruleset.StreamInfo) ruleset.MatchResult {
	return ruleset.MatchResult{Action: ruleset.ActionMaybe}
}

func benchmarkUDPManager(b *testing.B, churn bool) {
	node, err := snowflake.NewNode(0)
	if err != nil {
		b.Fatalf("create node: %v", err)
	}
	factory := &udpStreamFactory{WorkerID: 0, Logger: noopTestLogger{}, Node: node, Ruleset: emptyRuleset{}}
	mgr, err := newUDPStreamManager(factory, 200000, &statsCounters{})
	if err != nil {
		b.Fatalf("new manager: %v", err)
	}

	const flowCount = 20000
	tuples := make([]udpTupleKey, flowCount)
	payloads := make([][]byte, flowCount)
	for i := 0; i < flowCount; i++ {
		a := byte(i >> 8)
		c := byte(i)
		var t udpTupleKey
		t.AIP = [16]byte{10, a, 0, c}
		t.ALen = 4
		t.BIP = [16]byte{172, 16, a, c}
		t.BLen = 4
		t.APort = 1024 + uint16(i%20000)
		t.BPort = 20000 + uint16((i*7)%20000)
		tuples[i] = t
		payloads[i] = []byte{0x01, 0x00, 0x00, 0x00}
	}

	ctx := &udpContext{Verdict: udpVerdictAccept}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		idx := i % flowCount
		streamID := uint32(idx + 1)
		if churn {
			streamID = uint32((i % flowCount) + 1 + ((i / flowCount) * flowCount))
		}
		ctx.Verdict = udpVerdictAccept
		ctx.Packet = nil
		mgr.MatchWithContext(streamID, tuples[idx], false, payloads[idx], ctx)
	}
}

func BenchmarkUDPManagerMatchStableStreamID(b *testing.B) {
	benchmarkUDPManager(b, false)
}

func BenchmarkUDPManagerMatchStreamIDChurn(b *testing.B) {
	benchmarkUDPManager(b, true)
}

func BenchmarkLegacyUDPFallbackScanChurn(b *testing.B) {
	const flowCount = 5000
	tuples := make([]udpTupleKey, flowCount)
	for i := 0; i < flowCount; i++ {
		a := byte(i >> 8)
		c := byte(i)
		var t udpTupleKey
		t.AIP = [16]byte{10, a, 0, c}
		t.ALen = 4
		t.BIP = [16]byte{172, 16, a, c}
		t.BLen = 4
		t.APort = 1024 + uint16(i%20000)
		t.BPort = 20000 + uint16((i*7)%20000)
		tuples[i] = t
	}

	streams := make(map[uint32]*legacyUDPStreamValue, flowCount)
	keys := make([]uint32, 0, flowCount)
	for i := 0; i < flowCount; i++ {
		streamID := uint32(i + 1)
		streams[streamID] = &legacyUDPStreamValue{Tuple: tuples[i]}
		keys = append(keys, streamID)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		idx := i % flowCount
		streamID := uint32((i % flowCount) + 1 + ((i / flowCount) * flowCount))
		if _, ok := streams[streamID]; ok {
			continue
		}
		tuple := tuples[idx]
		revTuple := reverseTuple(tuple)
		for _, k := range keys {
			v, ok := streams[k]
			if !ok || v == nil {
				continue
			}
			if v.Tuple == tuple || v.Tuple == revTuple {
				delete(streams, k)
				streams[streamID] = v
				break
			}
		}
	}
}