// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Tetragon

//go:build !windows

// Package perfring provides utilities to do tests using the perf ringbuffer directly
package perfring

// NB(kkourt): Function(t *testing.T, ctx context.Context) is the reasonable
// thing to do here even if revive complains.
//revive:disable:context-as-argument

import (
	"context"
	"fmt"
	"path/filepath"
	"sync"
	"testing"

	"github.com/cilium/ebpf"
	"github.com/cilium/ebpf/perf"
	"github.com/cilium/ebpf/ringbuf"

	"github.com/cilium/tetragon/pkg/bpf"
	cfg "github.com/cilium/tetragon/pkg/config"
	testapi "github.com/cilium/tetragon/pkg/grpc/test"
	"github.com/cilium/tetragon/pkg/logger"
	"github.com/cilium/tetragon/pkg/observer"
	"github.com/cilium/tetragon/pkg/option"
	"github.com/cilium/tetragon/pkg/reader/notify"
	"github.com/cilium/tetragon/pkg/sensors/base"
	testsensor "github.com/cilium/tetragon/pkg/sensors/test"
	"github.com/cilium/tetragon/pkg/testutils"
)

// EventFn is the type of function called by ProcessEvents for each event.
type EventFn func(ev notify.Message) error

func loopEvents(events []observer.Event, eventFn EventFn, complChecker *testsensor.CompletionChecker) error {
	for _, ev := range events {
		switch xev := ev.(type) {
		case *testapi.MsgTestEventUnix:
			cpu := xev.Msg.Arg0
			complChecker.Update(cpu)
		}

		if err := eventFn(ev); err != nil {
			return err
		}
	}

	return nil
}

// ProcessEvents will open the perf ringbuffer and process events.
//
// It will complete, whenever it sees a number of MsgTestEventUnix on _all_
// cpus or when the ctx is done.  Hence, callers need to load the test sensor
// before they call it or use the ctx to cancel its operation.
//
// Example code to load the test sensor:
//
//	import testsensor "github.com/cilium/tetragon/pkg/sensors/test"
//	import tus "github.com/cilium/tetragon/pkg/testutils/sensors"
//	tus.LoadSensor(t, testsensor.GetTestSensor())
//
// If the test sensor is loaded, users can use TestCheckerMarkEnd to generate
// the appropriate MsgTestEventUnix.
func ProcessEvents(t *testing.T, ctx context.Context, eventFn EventFn, wgStarted *sync.WaitGroup) {
	pinOpts := ebpf.LoadPinOptions{}
	config := bpf.DefaultPerfEventConfig()

	var ringBufMapName string
	var ringBufReader *ringbuf.Reader
	var perfReader *perf.Reader
	useBPFRingBuffer := cfg.EnableV511Progs() && !option.Config.UsePerfRingBuffer

	if useBPFRingBuffer {
		ringBufMapName = filepath.Join(bpf.MapPrefixPath(), base.RingBufMapName)
	} else {
		ringBufMapName = config.MapName
	}

	ringBufMap, err := ebpf.LoadPinnedMap(ringBufMapName, &pinOpts)
	if err != nil {
		t.Fatalf("opening pinned map '%s' failed: %v", config.MapName, err)
	}
	defer ringBufMap.Close()

	if useBPFRingBuffer {
		ringBufReader, err = ringbuf.NewReader(ringBufMap)
	} else {
		perfReader, err = perf.NewReader(ringBufMap, 65535)
	}
	if err != nil {
		t.Fatalf("creating perf array reader failed: %v", err)
	}

	wgStarted.Done()

	errChan := make(chan error)
	defer close(errChan)

	complChan := make(chan bool)
	defer close(complChan)

	var wg sync.WaitGroup
	wg.Add(1)
	defer wg.Wait()

	if useBPFRingBuffer {
		// Service the BPF ring buffer.
		go func() {
			defer wg.Done()

			complChecker := testsensor.NewCompletionChecker()

			for ctx.Err() == nil {

				record, err := ringBufReader.Read()
				if err != nil {
					if ctx.Err() == nil {
						errChan <- fmt.Errorf("error reading ringbuf data: %w", err)
					}
					break
				}

				_, events, handlerErr := observer.HandlePerfData(record.RawSample)
				if handlerErr != nil {
					errChan <- fmt.Errorf("error handling ringbuf data: %w", handlerErr)
					break
				}
				err = loopEvents(events, eventFn, complChecker)
				if err != nil {
					errChan <- fmt.Errorf("error loop event function returned: %w", err)
					break
				}

				if complChecker.Done() {
					complChan <- true
					break
				}
			}
		}()
	} else {
		go func() {
			defer wg.Done()

			complChecker := testsensor.NewCompletionChecker()

			for ctx.Err() == nil {

				record, err := perfReader.Read()
				if err != nil {
					if ctx.Err() == nil {
						errChan <- fmt.Errorf("error reading perfring data: %w", err)
					}
					break
				}

				_, events, handlerErr := observer.HandlePerfData(record.RawSample)
				if handlerErr != nil {
					errChan <- fmt.Errorf("error handling perfring data: %w", handlerErr)
					break
				}
				err = loopEvents(events, eventFn, complChecker)
				if err != nil {
					errChan <- fmt.Errorf("error loop event function returned: %w", err)
					break
				}

				if complChecker.Done() {
					complChan <- true
					break
				}
			}
		}()
	}

	for {
		select {
		case err := <-errChan:
			t.Fatal(err)
		case <-complChan:
			if useBPFRingBuffer {
				ringBufReader.Close()
			} else {
				perfReader.Close()
			}
			return
		case <-ctx.Done():
			// Wait for context cancel.
			if useBPFRingBuffer {
				ringBufReader.Close()
			} else {
				perfReader.Close()
			}
			return
		}
	}
}

// RunTest is a convinience wrapper around ProcessEvents
// it will:
//   - run ProcessEvents on a different goroutinem using eventFn
//   - execute selfOperations provided by the user
//   - execute TestCheckerMarkEnd so that ProcessEvents returns
//   - wait for ProcessEvents to finish, and return
func RunTest(t *testing.T, ctx context.Context, selfOperations func(), eventFn EventFn) {
	var wgDone, wgStarted sync.WaitGroup
	wgDone.Add(1)
	wgStarted.Add(1)
	go func() {
		defer wgDone.Done()
		ProcessEvents(t, ctx, eventFn, &wgStarted)
	}()
	wgStarted.Wait()
	selfOperations()
	testsensor.TestCheckerMarkEnd(t)
	wgDone.Wait()
}

func FilterTestMessages(n notify.Message) bool {
	_, ok := n.(*testapi.MsgTestEventUnix)
	return !ok
}

// RunTestEvents: returns a list of events after running the test with selfOperations (see RunTest)
// MsgTestEventUnix events are filtered out.
func RunTestEvents(t *testing.T, ctx context.Context, selfOperations func()) []notify.Message {
	return RunTestEventReduce(t, ctx, selfOperations,
		FilterTestMessages,
		func(x notify.Message) notify.Message {
			return x
		},
		func(arr []notify.Message, m notify.Message) []notify.Message {
			return append(arr, m)
		},
	)
}

func RunTestEventReduce[K any, V any](
	t *testing.T,
	ctx context.Context,
	selfOperations func(),
	filterFn func(notify.Message) bool,
	mapFn func(notify.Message) K,
	reduceFn func(val V, key K) V,
) V {
	var ret V
	RunTest(t, ctx, selfOperations, func(ev notify.Message) error {
		if !filterFn(ev) {
			return nil
		}
		key := mapFn(ev)
		ret = reduceFn(ret, key)
		return nil
	})
	return ret
}

func RunTestEventReduceCount[K comparable](
	t *testing.T,
	ctx context.Context,
	selfOperations func(),
	filterFn func(notify.Message) bool,
	mapFn func(notify.Message) K,
) map[K]int {
	return RunTestEventReduce[K, map[K]int](
		t, ctx, selfOperations, filterFn, mapFn,
		func(v map[K]int, k K) map[K]int {
			if v == nil {
				v = make(map[K]int)
			}
			v[k]++
			return v
		},
	)
}

// similar to RunTest, but uses t.Run()
func RunSubTest(t *testing.T, ctx context.Context, name string, selfOperations func(t *testing.T), eventFn EventFn) bool {
	return t.Run(name, func(t *testing.T) {
		testutils.CaptureLog(t, logger.GetLogger())
		RunTest(t, ctx, func() { selfOperations(t) }, eventFn)
	})
}
