package tracer

import (
	"fmt"
	"github.com/coroot/coroot-node-agent/ebpftracer/tracer/ptrace"
	"github.com/coroot/coroot-node-agent/utils"
	klog "github.com/sirupsen/logrus"
	"math"
	"os"
	"runtime"
)

type AllocationDetails struct {
	StartAddr uint64
	EndAddr   uint64
	//NumCPU    uint64
}

func Allocate(pid int) (*AllocationDetails, error) {
	nCPU, err := utils.GetCPUCount()
	if err != nil {
		return nil, err
	}

	mapSize := uint64(os.Getpagesize() * nCPU * 8)

	addr, err := remoteAllocate(pid, mapSize)
	if err != nil {
		return nil, err
	}
	return &AllocationDetails{
		StartAddr: addr,
		EndAddr:   addr + mapSize,
		//NumCPU:    uint64(nCPU),
	}, nil
}

func remoteAllocate(pid int, mapSize uint64) (uint64, error) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	program, err := ptrace.NewTracedProgram(pid)
	if err != nil {
		return 0, err
	}

	defer func() {
		klog.Info("Detaching from process ", pid)
		err := program.Detach()
		if err != nil {
			klog.Error(err, "Failed to detach ptrace", "pid", pid)
		}
	}()

	if err := program.SetMemLockInfinity(); err != nil {
		klog.Error(err, "Failed to set memlock on process")
	} else {
		klog.Info("Set memlock on process successfully")
	}

	fd := -1
	addr, err := program.Mmap(mapSize, uint64(fd))
	if err != nil {
		return 0, err
	}
	if addr == math.MaxUint64 {
		// On success, mmap() returns a pointer to the mapped area.
		// On error, the value MAP_FAILED (that is, (void *) -1) is returned
		return 0, fmt.Errorf("mmap MAP_FAILED")
	}

	err = program.Madvise(addr, mapSize)
	if err != nil {
		return 0, err
	}

	err = program.Mlock(addr, mapSize)
	if err != nil {
		return 0, err
	}

	return addr, nil
}