package containers

import (
	"debug/dwarf"
	"debug/elf"
	debugelf "debug/elf"
	"fmt"
	"github.com/coroot/coroot-node-agent/common"
	"github.com/coroot/coroot-node-agent/ebpftracer"
	"github.com/coroot/coroot-node-agent/ebpftracer/tracer"
	tracerelf "github.com/coroot/coroot-node-agent/ebpftracer/tracer"
	"github.com/coroot/coroot-node-agent/proc"
	"github.com/coroot/coroot-node-agent/utils"
	. "github.com/coroot/coroot-node-agent/utils/modelse"
	klog "github.com/sirupsen/logrus"
	"golang.org/x/arch/arm64/arm64asm"
	"golang.org/x/arch/x86/x86asm"
	"io"
	"os"
	"regexp"
	"sort"
)

type uprobesDef struct {
	Name       string
	Offset     uint64
	EntAddress uint64
	RetAddress uint64
}

func (c *Container) AttachStack(tracer *ebpftracer.Tracer, pid uint32) error {
	//// 禁用stack
	//if tracer.DisableStackTracing() {
	//	klog.Warnf("StackTrace tracing is disabled")
	//	return nil
	//}

	if common.IsOpenFilter() && !common.IsFilterPid(pid) {
		klog.Warnf("StackTrace %d tracing is filter", pid)
		return nil
	}

	codeType := c.GetCodeTypeFromCache(pid)
	if codeType.IsUnknownCode() {
		klog.Warnf("StackTrace %d tracing is IsUnknownCode", pid)
		return nil
	}

	p := c.processes[pid]
	if p == nil {
		return fmt.Errorf("unknown process %d", pid)
	}

	if p.stackAttachOnce {
		return nil
	}

	p.stackAttachOnce = true

	switch codeType {
	case CodeTypeJava:
		return c.jvmStackTrace(tracer, pid)
	default:
		return c.stackTrace(tracer, pid)
	}
}

func (c *Container) stackTrace(tracer *ebpftracer.Tracer, pid uint32) error {
	p := c.processes[pid]
	if p == nil {
		return fmt.Errorf("unknown process %d", pid)
	}

	//if p.stackAttachOnce {
	//	return nil
	//}

	binType := "dotnet"
	MatchString := ".*HandleFunc|.*main.*|testfun.*|.*serverHandler.*|.*ServeHTTP.*"
	dbgpath := ""

	WHITE_LIST := os.Getenv("WHITE_LIST")
	BIN_TYPE := os.Getenv("BIN_TYPE")
	DBG_PATH := os.Getenv("DBG_PATH")
	if WHITE_LIST != "" {
		MatchString = WHITE_LIST
	}

	if DBG_PATH != "" {
		dbgpath = DBG_PATH
	}

	if BIN_TYPE != "" {
		binType = BIN_TYPE
	}

	klog.Infoln("[stack] UprobesMatchString:::init", MatchString)
	path := proc.Path(uint32(pid), "exe")

	var err error
	if dbgpath != "" {
		c.Uprobes, err = c.getJavaAOTUprobes(binType, path, dbgpath, MatchString)
	} else {
		c.Uprobes, err = c.getUprobes(path, MatchString)
	}

	if err != nil {
		return err
	}

	c.UprobesMap = map[string]tracerelf.Uprobe{}
	klog.Infoln("[stack] UprobesMap start")
	for _, up := range c.Uprobes {
		klog.Debugf("[stack] UprobesMap %s %d %d", up.Funcname, up.Address, up.AbsOffset)
		c.UprobesMap[fmt.Sprintf("%s-%s", up.Funcname, up.Address+up.AbsOffset)] = up
	}
	//codeType := c.GetCodeTypeFromCache(pid)
	//tracer.InitKProcInfo(pid, c.instanceID, uint16(codeType))
	p.stackUprobes = append(p.stackUprobes, tracer.AttachStackUprobes(path, c.Uprobes)...)
	p.stackAttachOnce = true

	return nil
}

func (c *Container) jvmStackTrace(tracer *ebpftracer.Tracer, pid uint32) error {
	p := c.processes[pid]
	// check version
	libjavaso, err := utils.GetSoPath(pid, "libjava.so", c.getRootfs())
	if err != nil {
		p.versionFailed = true
		klog.WithError(err).Errorf("[jvmStackTrace] Failed get so path")
		return err
	}
	libjvmso, err := utils.GetSoPath(pid, "libjvm.so", c.getRootfs())
	if err != nil {
		klog.WithError(err).Errorf("[jvmStackTrace] Failed get so path")
		return err
	}
	v, err := ebpftracer.GetJvmVersion(libjavaso, libjvmso)
	if err != nil {
		p.versionFailed = true
		klog.WithError(err).Errorf("[jvmStackTrace] Failed get Java version")
		return err
	}
	c.AppInfo.Version = v
	major, minor, patch, err := ebpftracer.ParseVersion(v)
	klog.Infof("[jvmStackTrace] version: %s (Major: %d, Minor: %d, Patch: %d)", v, major, minor, patch)
	if major != 1 || minor != 8 {
		p.versionFailed = true
		return fmt.Errorf("[jvmStackTrace] Unsupported Java version")
	}

	err = tracer.JattachJvm(pid, c.AppInfo, c.WhiteSettingInfo.WhiteStackSettingInfo.WhiteList, c.WhiteSettingInfo.WhiteStackSettingInfo.BlackList, c.getRootfs())
	if err != nil {
		p.stackStatus.JattachFailure()
		return err
	} else {
		p.stackStatus.JattachSuccess()
	}

	jvmStackProbes, err := tracer.AttachJVMStackUprobes(pid, c.AppInfo, c.getRootfs())

	if err != nil {
		p.stackStatus.StackUprobesFailure()
		klog.WithError(err).Errorf("[jvmStackTrace] Failed attach jvm stack.")
		return err
	}

	p.stackUprobes = append(p.stackUprobes, jvmStackProbes...)
	p.stackStatus.StackUprobesSuccess()

	return nil
}

func (c *Container) getJavaAOTUprobes(binType, path string, dbgpath string, MatchString string) ([]tracer.Uprobe, error) {
	uprobes := []tracer.Uprobe{}

	elfFile, err := elf.Open(path)
	if err != nil {
		return nil, err
	}
	funSection := ".text"

	if binType == "dotnet" {
		funSection = "__managedcode"
	}

	textSection := elfFile.Section(funSection)
	if textSection == nil {
		//fmt.Println("no text section", nil)
		return nil, nil
	}
	textSectionData, err := textSection.Data()
	if err != nil {
		//fmt.Println("failed to read text section", err)
		return nil, err
	}
	textSectionLen := uint64(len(textSectionData) - 1)

	dwarfFile, err := elf.Open(dbgpath)

	if err != nil {
		return nil, err
	}

	dwarfData, err := dwarfFile.DWARF()
	if err != nil {
		return nil, err
	}

	entryReader := dwarfData.Reader()

	// var targetAddress uint64

	listEntry := make(map[dwarf.Offset]uprobesDef)
	SpecListEntry := []dwarf.Entry{}

	for {
		entry, err := entryReader.Next()
		if err == io.EOF {
			// We've reached the end of DWARF entries
			break
		}
		if err != nil {
			//log.Fatalf("Error reading entry: %v", err)
			klog.Fatalf("Error reading entry: %v", err)
		}
		if entry == nil {
			//log.Println("Warning: a nil entry was returned with no error")
			klog.Warn("Warning: a nil entry was returned with no error")
			break
		}
		if entry.Tag == dwarf.TagSubprogram {
			// fmt.Printf("entry address: %x, %d\n", entry.Offset, entry.Children)
			funName, _ := entry.Val(dwarf.AttrName).(string)
			found, _ := regexp.MatchString(MatchString, funName)
			if found {
				entAddress, _ := entry.Val(dwarf.AttrLowpc).(uint64)
				retAddress, _ := entry.Val(dwarf.AttrHighpc).(uint64)
				// fmt.Printf("Function %s address: %x, %x\n", funName, address, entry.Offset)
				uprobes := uprobesDef{}
				uprobes.EntAddress = entAddress
				uprobes.RetAddress = retAddress
				uprobes.Offset = uint64(entry.Offset)
				uprobes.Name = funName

				listEntry[entry.Offset] = uprobes
			}

			specAddr, _ := entry.Val(dwarf.AttrSpecification).(dwarf.Offset)
			lowpc := entry.Val(dwarf.AttrLowpc)
			if lowpc != nil && specAddr > 0 && lowpc.(uint64) > 0 {
				// fmt.Printf("AttrSpecification address: %x, %x\n", specAddr, entry.Offset)
				SpecListEntry = append(SpecListEntry, *entry)
			}
		}
	}

	for _, v := range SpecListEntry {
		specAddr, _ := v.Val(dwarf.AttrSpecification).(dwarf.Offset)
		// fmt.Printf("SpecListEntrySpecListEntrySpecListEntry Attach Function: %x\n", specAddr)
		_, ok := listEntry[specAddr]
		if ok {
			vv := listEntry[specAddr]
			entAddr := v.Val(dwarf.AttrLowpc)
			if entAddr != nil {
				vv.EntAddress = entAddr.(uint64)
			}
			retAddr := v.Val(dwarf.AttrHighpc)
			if retAddr != nil {
				switch retAddr.(type) {
				case uint64:
					vv.RetAddress = uint64(retAddr.(uint64))
				case int64:
					vv.RetAddress = uint64(retAddr.(int64))
				default:
					//fmt.Println("Unknown type")
				}
			}
			listEntry[specAddr] = vv
		}
	}

	for _, v := range listEntry {
		//fmt.Printf("Need Attach Function %s address: %x, %x\n", v.Name, v.EntAddress, v.RetAddress)

		sStart := v.EntAddress - textSection.Addr
		sSize := v.RetAddress
		if v.RetAddress > v.EntAddress {
			sSize = v.RetAddress - v.EntAddress
		}
		sEnd := sStart + sSize
		if sEnd > textSectionLen {
			continue
		}
		sBytes := textSectionData[sStart:sEnd]
		rbpOffsets := getRbpEnterOffsets(elfFile.Machine, sBytes)
		returnOffsets := getReturnOffsets(elfFile.Machine, sBytes)

		if rbpOffsets != 0 {
			uprobes = append(uprobes, tracer.Uprobe{
				Funcname:  v.Name,               // 函数名
				Location:  tracer.AtDotNetEntry, // 入口
				Address:   v.EntAddress,         // 函数地址
				AbsOffset: uint64(rbpOffsets),   // 函数相对 ELF 偏移
				RelOffset: 0,                    // 函数真实偏移
			})
		} else {
			// 函数入口加入待 attach 列表
			uprobes = append(uprobes, tracer.Uprobe{
				Funcname:  v.Name,         // 函数名
				Location:  tracer.AtEntry, // 入口
				Address:   v.EntAddress,   // 函数地址
				AbsOffset: 0,              // 函数相对 ELF 偏移
				RelOffset: 0,              // 函数真实偏移
			})
		}

		for _, offset := range returnOffsets {
			uprobes = append(uprobes, tracer.Uprobe{
				Funcname:  v.Name,
				Location:  tracer.AtRet,
				Address:   v.EntAddress,
				AbsOffset: uint64(offset),
				RelOffset: 0,
			})
		}
	}
	return uprobes, nil
}

func (c *Container) getUprobes(path string, MatchString string) ([]tracer.Uprobe, error) {
	uprobes := []tracer.Uprobe{}

	binFile, err := os.Open(path)
	if err != nil {
		return nil, err
	}

	// cache := map[string]interface{}{}
	// 解析 elf 文件
	elfFile, err := debugelf.NewFile(binFile)
	if err != nil {
		return nil, err
	}
	// 获取所有符号表
	symbols, err := elfFile.Symbols()
	if err != nil {
		return nil, err
	}
	sort.Slice(symbols, func(i, j int) bool { return symbols[i].Value < symbols[j].Value })

	c.Symbols = symbols
	// 符号表组装成键值 map，方便使用
	symnames := map[string]debugelf.Symbol{}
	for _, symbol := range symbols {
		klog.Debugf("[stack] %v %v", symbol.Name, symbol)
		symnames[symbol.Name] = symbol
	}

	textSection := elfFile.Section(".text")
	if textSection == nil {
		klog.Infoln("[stack] no text section")
		return nil, nil
	}
	textSectionData, err := textSection.Data()
	if err != nil {
		klog.WithError(err).Errorf("[stack] Failed to read text section")
		return nil, nil
	}
	textSectionLen := uint64(len(textSectionData) - 1)

	//  遍历符号表
	for _, symbol := range symbols {
		if debugelf.ST_TYPE(symbol.Info) != debugelf.STT_FUNC {
			continue
		}
		// fmt.Println("Hello FunName: ", symbol.Name)
		// 使用正则表达式匹配函数白名单列表
		found, err := regexp.MatchString(MatchString, symbol.Name)
		// found, err := regexp.MatchString("main.*", symbol.Name)

		if err != nil {
			klog.WithError(err).Errorln("[stack] found error")
			return nil, err
		}

		if found {
			// 匹配到了加入 attachFuncs 列表
			klog.Debugf("[stack] Fuck This: %s, %x", symbol.Name, symbol.Value)
			// attachFuncs = append(attachFuncs, symbol.Name)
			// 根据函数名拿到当前函数的符号结构体
			sym := symnames[symbol.Name]
			//if err != nil {
			//	klog.WithError(err).Errorf("symnames[symbol.Name] %s", symbol.Name)
			//	return nil, err
			//}

			address := sym.Value
			for _, p := range elfFile.Progs {
				if p.Type != elf.PT_LOAD || (p.Flags&elf.PF_X) == 0 {
					continue
				}

				if p.Vaddr <= sym.Value && sym.Value < (p.Vaddr+p.Memsz) {
					address = sym.Value - p.Vaddr + p.Off
					break
				}
			}

			// 函数入口加入待 attach 列表
			uprobes = append(uprobes, tracer.Uprobe{
				Funcname:  symbol.Name,    // 函数名
				Location:  tracer.AtEntry, // 入口
				Address:   address,        // 函数地址
				AbsOffset: 0,              // 函数相对 ELF 偏移
				RelOffset: 0,              // 函数真实偏移
				Wanted:    true,
			})

			sStart := sym.Value - textSection.Addr
			sEnd := sStart + sym.Size
			if sEnd > textSectionLen {
				continue
			}
			sBytes := textSectionData[sStart:sEnd]
			returnOffsets := getReturnOffsets(elfFile.Machine, sBytes)

			for _, offset := range returnOffsets {
				uprobes = append(uprobes, tracer.Uprobe{
					Funcname:  symbol.Name,
					Location:  tracer.AtRet,
					Address:   address,
					AbsOffset: uint64(offset),
					RelOffset: 0,
				})
			}
		}
	}
	return uprobes, nil
}

func getRbpEnterOffsets(machine elf.Machine, instructions []byte) int {
	switch machine {
	case elf.EM_X86_64:
		for i := 0; i < len(instructions); {
			ins, err := x86asm.Decode(instructions[i:], 64)
			if err == nil && ins.Op == x86asm.LEA && ins.Args[0].String() == "RBP" {
				klog.Infof("[stack] getRbpEnterOffsets: %v, %s, %s", ins, ins.Args[0].String(), ins.Args[1].String())
				return i
			}
			i += ins.Len
		}
	case elf.EM_AARCH64:
		for i := 0; i < len(instructions); {
			ins, err := arm64asm.Decode(instructions[i:])
			if err == nil && ins.Op == arm64asm.RET {
				return i
			}
			i += 4
		}
	}
	return 0
}

func getReturnOffsets(machine elf.Machine, instructions []byte) []int {
	var res []int
	switch machine {
	case elf.EM_X86_64:
		for i := 0; i < len(instructions); {
			ins, err := x86asm.Decode(instructions[i:], 64)
			if err == nil && ins.Op == x86asm.RET {
				res = append(res, i)
			}
			i += ins.Len
		}
	case elf.EM_AARCH64:
		for i := 0; i < len(instructions); {
			ins, err := arm64asm.Decode(instructions[i:])
			if err == nil && ins.Op == arm64asm.RET {
				res = append(res, i)
			}
			i += 4
		}
	}
	return res
}