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

//go:build !windows

package btf

import (
	"fmt"
	"math"
	"os"
	"path"
	"reflect"
	"regexp"
	"strconv"
	"strings"

	"github.com/cilium/ebpf/btf"

	api "github.com/cilium/tetragon/pkg/api/tracingapi"
	"github.com/cilium/tetragon/pkg/defaults"
	"github.com/cilium/tetragon/pkg/kernels"
	"github.com/cilium/tetragon/pkg/ksyms"
	"github.com/cilium/tetragon/pkg/logger"
	"github.com/cilium/tetragon/pkg/option"
)

var (
	btfFile string
)

func btfFileExists(file string) error {
	_, err := os.Stat(file)
	return err
}

func observerFindBTF(lib, btf string) (string, error) {
	if btf == "" {
		// Alternative to auto-discovery and/or command line argument we
		// can also set via environment variable.
		tetragonBTFEnv := os.Getenv("TETRAGON_BTF")
		if tetragonBTFEnv != "" {
			if _, err := os.Stat(tetragonBTFEnv); err != nil {
				return btf, err
			}
			return tetragonBTFEnv, nil
		}

		_, kernelVersion, err := kernels.GetKernelVersion(option.Config.KernelVersion, option.Config.ProcFS)
		if err != nil {
			return btf, err
		}

		// Preference of BTF files, first search for kernel exposed BTF, then
		// check for vmlinux- hubble metadata, and finally if all those are missing
		// search the lib directory for a btf file.
		if _, err := os.Stat(defaults.DefaultBTFFile); err == nil {
			logger.GetLogger().Info("BTF discovery: default kernel btf file found", "btf-file", defaults.DefaultBTFFile)
			return defaults.DefaultBTFFile, nil
		}
		logger.GetLogger().Info("BTF discovery: default kernel btf file does not exist", "btf-file", defaults.DefaultBTFFile)

		runFile := path.Join(lib, "metadata", "vmlinux-"+kernelVersion)
		if _, err := os.Stat(runFile); err == nil {
			logger.GetLogger().Info("BTF discovery: candidate btf file found", "btf-file", runFile)
			return runFile, nil
		}
		logger.GetLogger().Info("BTF discovery: candidate btf file does not exist", "btf-file", runFile)

		runFile = path.Join(lib, "btf")
		if _, err := os.Stat(runFile); err == nil {
			logger.GetLogger().Info("BTF discovery: candidate btf file found", "btf-file", runFile)
			return runFile, nil
		}
		logger.GetLogger().Info("BTF discovery: candidate btf file does not exist", "btf-file", runFile)

		return btf, fmt.Errorf("kernel version %q BTF search failed kernel is not included in supported list. Please check Tetragon requirements documentation, then use --btf option to specify BTF path and/or '--kernel' to specify kernel version", kernelVersion)
	}
	if err := btfFileExists(btf); err != nil {
		return btf, fmt.Errorf("user specified BTF does not exist: %w", err)
	}
	logger.GetLogger().Info("BTF file: user specified btf file found", "btf-file", btf)
	return btf, nil
}

type Spec = btf.Spec

func NewBTF() (*btf.Spec, error) {
	return btf.LoadSpec(btfFile)
}

func InitCachedBTF(lib, btf string) error {
	var err error

	// Find BTF metdaata and populate btf opaqu object
	btfFile, err = observerFindBTF(lib, btf)
	if err != nil {
		return fmt.Errorf("tetragon, aborting kernel autodiscovery failed: %w", err)
	}
	return err
}

func GetCachedBTFFile() string {
	return btfFile
}

func FindBTFStruct(name string) (*btf.Struct, error) {
	var ty *btf.Struct

	spec, err := NewBTF()
	if err != nil {
		return nil, err
	}

	err = firstTypeByName(spec, name, &ty)
	return ty, err
}

// firstStructTypeByName mimics spec.TypeByName(), but returns first match found.
func firstTypeByName(spec *btf.Spec, name string, typ any) error {
	typeInterface := reflect.TypeFor[btf.Type]()

	// typ may be **T or *Type
	typValue := reflect.ValueOf(typ)
	if typValue.Kind() != reflect.Ptr {
		return fmt.Errorf("%T is not a pointer", typ)
	}

	typPtr := typValue.Elem()
	if !typPtr.CanSet() {
		return fmt.Errorf("%T cannot be set", typ)
	}

	wanted := typPtr.Type()
	if wanted == typeInterface {
		// This is *Type. Unwrap the value's type.
		wanted = typPtr.Elem().Type()
	}

	if !wanted.AssignableTo(typeInterface) {
		return fmt.Errorf("%T does not satisfy Type interface", typ)
	}

	types, err := spec.AnyTypesByName(name)
	if err != nil {
		return err
	}

	for _, typ := range types {
		if reflect.TypeOf(typ) != wanted {
			continue
		}
		typPtr.Set(reflect.ValueOf(typ))
		return nil
	}
	return btf.ErrNotFound
}

func FindBTFFuncParamFromHook(hook string, argIndex int) (*btf.FuncParam, error) {
	// If the hook is part of a kernel module, load its BTF file
	if ks, err := ksyms.KernelSymbols(); err == nil {
		if kmod, err := ks.GetKmod(hook); err == nil {
			spec, err := btf.LoadKernelModuleSpec(kmod)
			if err != nil {
				return nil, err
			}
			return findBTFFuncParamFromHookWithSpec(spec, hook, argIndex)
		}
	}
	// In case of failure, default to kernel BTF
	spec, err := NewBTF()
	if err != nil {
		return nil, err
	}
	return findBTFFuncParamFromHookWithSpec(spec, hook, argIndex)
}

func findBTFFuncParamFromHookWithSpec(spec *btf.Spec, hook string, argIndex int) (*btf.FuncParam, error) {
	var hookFn *btf.Func

	if err := spec.TypeByName(hook, &hookFn); err != nil {
		if strings.HasPrefix(hook, "bpf_lsm_") {
			return nil, fmt.Errorf("failed to find BTF type for hook %q: %w."+
				"Please check if the hook exists or if your kernel supports BTF for lsm hooks."+
				"As an alternative, consider switching to kprobes. ", hook, err)
		}
		return nil, fmt.Errorf("failed to find BTF type for hook %q: %w", hook, err)
	}

	btfHookProto, isBTFFuncProto := hookFn.Type.(*btf.FuncProto)
	if !isBTFFuncProto {
		return nil, fmt.Errorf("hook %q has no BTF type FuncProto", hook)
	}
	paramLen := len(btfHookProto.Params)
	if argIndex > paramLen-1 {
		parameter := "parameter"
		if paramLen > 1 {
			parameter += "s"
		}
		return nil, fmt.Errorf("index %d is out of range. The hook only have %d %q", argIndex, paramLen, parameter)
	}
	return &btfHookProto.Params[argIndex], nil
}

func ResolveNestedTypes(ty btf.Type) btf.Type {
	switch t := ty.(type) {
	case *btf.Restrict:
		return ResolveNestedTypes(t.Type)
	case *btf.Volatile:
		return ResolveNestedTypes(t.Type)
	case *btf.Const:
		return ResolveNestedTypes(t.Type)
	case *btf.Typedef:
		return ResolveNestedTypes(t.Type)
	}
	return ty
}

func parseArrayIdxStr(s string) (uint32, error) {
	re := regexp.MustCompile(`^\[(\d+)\]$`)

	matches := re.FindStringSubmatch(s)
	if len(matches) < 2 {
		return 0, fmt.Errorf("invalid format %q (must be: \"[value]\")", s)
	}

	n, err := strconv.Atoi(matches[1])
	if err != nil {
		return 0, fmt.Errorf("invalid value %q: %w", matches[1], err)
	}

	if n < 0 || n > math.MaxUint32 {
		return 0, fmt.Errorf("value %d out of range for uint32", n)
	}

	return uint32(n), nil
}

func getSizeofType(t btf.Type) uint32 {
	ret, _ := btf.Sizeof(t)
	return uint32(ret)
}

// ResolveBTFPath function recursively search in a btf structure in order to
// found a specific path until it reach the target or fail.

// The function also search in embedded anonymous structures or unions to cover as
// much use cases as possible. For instance, mm_struct have 2 fields, anonymous
// struct and another type. But you are still able to look into the anonymous
// struct by specifying a path like "mm.pgd.pgd".

// @btfArgs: dest array for storing btf informations to reach the target on the
// bpf side.
// @currentType: The current type being proccessed, starts with root type.
// @pathToFound: The string representation of the path to reach in the structures.
// @i: The current depth, until last element of pathToFound.

// Return: The last type found matching the path, or error.
func ResolveBTFPath(
	btfArgs *[api.MaxBTFArgDepth]api.ConfigBTFArg,
	currentType btf.Type,
	pathToFound []string,
	i int,
) (*btf.Type, error) {
	switch t := currentType.(type) {
	case *btf.Struct:
		return processMembers(btfArgs, currentType, t.Members, pathToFound, i)
	case *btf.Union:
		return processMembers(btfArgs, currentType, t.Members, pathToFound, i)
	case *btf.Pointer:
		if i > 0 {
			// To avoid adding an extra BTFArg entry only for doing this
			// dereferencing, we mark the previous element as pointer.
			btfArgs[i-1].IsPointer = uint16(1)
		}
		if idx, err := parseArrayIdxStr(pathToFound[i]); err == nil {
			// To stay ahead on the dereferecing, we mark the current btfArg as pointer
			btfArgs[i].IsPointer = uint16(1)
			return processArray(btfArgs, ResolveNestedTypes(t.Target), pathToFound, i, idx)
		}
		return ResolveBTFPath(btfArgs, ResolveNestedTypes(t.Target), pathToFound, i)
	case *btf.Array:
		idx, err := parseArrayIdxStr(pathToFound[i])
		if err != nil {
			return nil, fmt.Errorf("fail parsing array index : %w", err)
		}
		if idx >= t.Nelems {
			return nil, fmt.Errorf("array index out of bound. Nelems=%d, got=%d", t.Nelems, idx)
		}
		return processArray(btfArgs, ResolveNestedTypes(t.Type), pathToFound, i, idx)
	default:
		ty := currentType.TypeName()
		if len(ty) == 0 {
			ty = reflect.TypeOf(currentType).String()
		}
		currentPath := pathToFound[i]
		if i > 0 {
			currentPath = pathToFound[i-1]
		}
		return nil, fmt.Errorf("unexpected type : %q has type %q", currentPath, ty)
	}
}

func processMembers(
	btfArgs *[api.MaxBTFArgDepth]api.ConfigBTFArg,
	currentType btf.Type,
	members []btf.Member,
	pathToFound []string,
	i int,
) (*btf.Type, error) {
	var lastError error
	memberWasFound := false
	for _, member := range members {
		if len(member.Name) == 0 { // If anonymous struct, fallthrough
			btfArgs[i].Offset = member.Offset.Bytes()
			btfArgs[i].IsInitialized = uint16(1)
			lastTy, err := ResolveBTFPath(btfArgs, ResolveNestedTypes(member.Type), pathToFound, i)
			if err != nil {
				if lastError != nil {
					idx := i + 1
					// If the error raised originates from a depth greater than the current one, we stop the search.
					if idx < len(pathToFound) && strings.Contains(lastError.Error(), pathToFound[idx]) {
						break
					}
				}
				lastError = err
				continue
			}
			return lastTy, nil
		}
		if member.Name == pathToFound[i] {
			memberWasFound = true
			btfArgs[i].Offset = member.Offset.Bytes()
			btfArgs[i].IsInitialized = uint16(1)
			isNotLastChild := i < len(pathToFound)-1 && i < api.MaxBTFArgDepth
			if isNotLastChild {
				return ResolveBTFPath(btfArgs, ResolveNestedTypes(member.Type), pathToFound, i+1)
			}
			currentType = ResolveNestedTypes(member.Type)
			break
		}
	}
	if !memberWasFound {
		if lastError != nil {
			return nil, lastError
		}
		return nil, fmt.Errorf(
			"attribute %q not found in structure %q found %v",
			pathToFound[i],
			currentType.TypeName(),
			members,
		)
	}
	switch t := currentType.(type) {
	case *btf.Pointer:
		btfArgs[i].IsPointer = uint16(1)
		currentType = t.Target
	case *btf.Int, *btf.Enum:
		btfArgs[i].IsPointer = uint16(1)
	}
	return &currentType, nil
}

func processArray(
	btfArgs *[api.MaxBTFArgDepth]api.ConfigBTFArg,
	targetType btf.Type,
	pathToFound []string,
	i int,
	idx uint32,
) (*btf.Type, error) {
	btfArgs[i].IsInitialized = uint16(1)
	btfArgs[i].Offset = getSizeofType(targetType) * idx
	if len(pathToFound) > i+1 {
		return ResolveBTFPath(btfArgs, targetType, pathToFound, i+1)
	}
	return &targetType, nil
}
