package internal

import "C"

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"runtime/pprof"
	"syscall"
	"time"

	"ior/internal/flags"
	"ior/internal/tracepoints"

	bpf "github.com/aquasecurity/libbpfgo"
)

func attachTracepoints(flags flags.Flags, bpfModule *bpf.Module) error {
	for _, name := range tracepoints.List {
		if !flags.AttachTracepoint(name) {
			continue
		}
		fmt.Println("Attaching tracepoint", name)

		prog, err := bpfModule.GetProgram(fmt.Sprintf("handle_%s", name))
		if err != nil {
			return fmt.Errorf("Failed to get BPF program handle_%s: %v", name, err)
		}
		fmt.Println("Attached prog handle_", name)

		if _, err = prog.AttachTracepoint("syscalls", name); err != nil {
			// OK, older Kernel versions may not have this tracepoint!
			fmt.Printf("Failed to attach to %s tracepoint: %v, kernel version may be too old, skipping", name, err)
			continue
		}
		fmt.Println("Attached tracepoint ", name)
	}

	return nil
}

func Run(flags flags.Flags) {
	bpfModule, err := bpf.NewModuleFromFile("ior.bpf.o")
	if err != nil {
		panic(err)
	}
	defer bpfModule.Close()

	if err := flags.ResizeBPFMaps(bpfModule); err != nil {
		panic(err)
	}

	if err := flags.SetBPF(bpfModule); err != nil {
		panic(err)
	}

	if err := bpfModule.BPFLoadObject(); err != nil {
		panic(err)
	}

	if err := attachTracepoints(flags, bpfModule); err != nil {
		panic(err)
	}

	// 4096 channel size, minimises event drops
	ch := make(chan []byte, 4096)
	rb, err := bpfModule.InitRingBuf("event_map", ch)
	if err != nil {
		panic(err)
	}
	rb.Poll(300)

	pprofDone := make(chan struct{})
	var cpuProfile, memProfile *os.File
	if flags.PprofEnable {
		if cpuProfile, err = os.Create("ior.cpuprofile"); err != nil {
			panic(err)
		}
		if memProfile, err = os.Create("ior.memprofile"); err != nil {
			panic(err)
		}
		pprof.StartCPUProfile(cpuProfile)
	} else {
		close(pprofDone)
	}

	loop := newEventLoop(flags)
	duration := time.Duration(flags.Duration) * time.Second
	fmt.Println("Probing for", duration)
	ctx, cancel := context.WithTimeout(context.Background(), duration)

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)

	go func() {
		<-c
		cancel()
	}()

	go func() {
		<-ctx.Done()
		fmt.Println(loop.stats())
		if flags.PprofEnable {
			fmt.Println("Stoppig profiling, writing ior.cpuprofile and ior.memprofile")
			pprof.StopCPUProfile()
			pprof.WriteHeapProfile(memProfile)
			close(pprofDone)
		}
	}()

	startTime := time.Now()
	loop.run(ctx, ch)
	totalDuration := time.Since(startTime)
	<-pprofDone
	fmt.Println("Good bye... (unloading BPF tracepoints will take a few seconds...) after", totalDuration)
}