package main

import (
	"encoding/json"
	"fmt"
	"os"
	"sort"
)

var count int
var starTgid int

func filtPids(pRawPidData *[]Process) {
	/* ATTENTION: 把map/slice直接传参是危险的
	 * 传递的是指针，不会引起大的复制开销，
	 * 但是map/slice在callee func内被修改**可能**导致内存更改
	 * 而这样的内存更改对caller function来说是不可见的，看到的还是原来的东西
	 * 这里由于参数几乎都是只读不写，因而用一下
	 */

	// 合并由多线程导致的重复记录，顺便按照pid升序
	// 多线程已经取消了，但保险起见还是留着
	merged := mergeProcess(pRawPidData)
	// 将Process按照tgid合并
	var tgidMap map[int]*tgidNode
	var rootfsPids []int
	tgidMap, starTgid, rootfsPids = getTgidNodes(merged)
	// 建树，helloTree
	buildTree(tgidMap, starTgid)
	// 对树上的进程做一些优化处理
	optimazePid(starTgid, rootfsPids)
}

func ProMerge(a, b Process) (res Process) {
	// 合并过程中会遇到什么问题？
	res.Star = false

	if a.StartTimestamp.IsZero() {
		res.StartTimestamp = b.StartTimestamp
	} else if b.StartTimestamp.IsZero() {
		res.StartTimestamp = a.StartTimestamp
	} else if a.StartTimestamp.Before(b.StartTimestamp) {
		res.StartTimestamp = a.StartTimestamp
	} else {
		res.StartTimestamp = b.StartTimestamp
	}

	res.Ppid = a.Ppid
	if a.ParentTgid == 0 {
		res.ParentTgid = b.ParentTgid
	} else {
		res.ParentTgid = a.ParentTgid
	}

	res.Pid = a.Pid
	if a.Tgid == 0 {
		res.Tgid = b.Tgid
	} else {
		res.Tgid = a.Tgid
	}

	if len(a.Args) == 0 {
		res.Args = b.Args
	} else {
		res.Args = a.Args
	}

	if a.Comm == "" {
		res.Comm = b.Comm
	} else {
		res.Comm = a.Comm
	}

	if a.RootFS == "" {
		res.RootFS = b.RootFS
	} else {
		res.RootFS = a.RootFS
	}

	if a.Cwd == "" {
		res.Cwd = b.Cwd
	} else {
		res.Cwd = a.Cwd
	}

	res.Execve = append(a.Execve, b.Execve...)
	res.Children = append(a.Children, b.Children...)

	var flag bool // 真a假b
	if a.ExitTimestamp.IsZero() {
		flag = false
	} else if b.ExitTimestamp.IsZero() {
		flag = true
	} else if a.ExitTimestamp.Before(b.ExitTimestamp) {
		flag = true
	} else {
		flag = false
	}

	if flag {
		res.ExitCode = a.ExitCode
		res.ExitSignal = a.ExitSignal
		res.ExitTimestamp = a.ExitTimestamp
	} else {
		res.ExitCode = b.ExitCode
		res.ExitSignal = b.ExitSignal
		res.ExitTimestamp = b.ExitTimestamp
	}

	return res
}

func mergeProcess(pRawPidData *[]Process) (merged []Process) {
	rawPidData := *pRawPidData
	// 合并由多线程导致的重复记录，顺便按照pid升序
	index := make(map[int]int)
	for _, process := range rawPidData {
		i, exists := index[process.Pid]
		if exists {
			// 已存在，合并
			merged[i] = ProMerge(merged[i], process)
		} else {
			// 不存在，直接添加
			merged = append(merged, process)
			index[process.Pid] = len(merged) - 1
		}
	}
	sort.Slice(merged, func(i, j int) bool {
		return merged[i].Pid < merged[j].Pid
	})
	return merged
}

func getTgidNodes(merged []Process) (tgidMap map[int]*tgidNode, starTgid int, rootfsPids []int) {
	// 合并出来的进程整理为tgidNode
	tgidMap = make(map[int]*tgidNode)
	findTgid = make(map[int]int) // pid --> tgid
	// var starTgid, rootFsPid int
	starTgid = -1
	// rootfsPid = -1
	rootfsPids = make([]int, 0)
	for _, val := range merged {
		if val.Star {
			starTgid = val.Tgid
		} else if val.RootFS != "" {
			rootfsPids = append(rootfsPids, val.Pid)
		}
		// 登记tgid
		findTgid[val.Pid] = val.Tgid
		nodeval, exists := tgidMap[val.Tgid]
		if exists {
			// 直接记录
			nodeval.Threads = append(nodeval.Threads, val)
			nodeval.FindPid[val.Pid] = len(nodeval.Threads) - 1
		} else {
			node := tgidNode{
				Tgid:      val.Tgid,
				FindPid:   make(map[int]int),
				Threads:   make([]Process, 0),
				ChildTgid: make([]int, 0),
			}
			node.Threads = append(node.Threads, val)
			node.FindPid[val.Pid] = 0
			tgidMap[val.Tgid] = &node
		}
	}
	return tgidMap, starTgid, rootfsPids
}

func buildTree(tgidMap map[int]*tgidNode, starTgid int) {
	// 从tgid==starTgid开始，构建树
	helloTree = make(map[int]*tgidNode) // 在树上的tgid节点，tgid --> *tgidNode
	var q Queue                         // 记录每一个整理好的结构体，bfs
	visited := make(map[int]bool)       // 哪些tgid已经访问过

	tmp, exists := tgidMap[starTgid]
	if !exists {
		return
	}

	// helloTree负责在遍历到该节点时记录
	// 队列仅负责搞明白哪些节点在树上
	// 因而所有添加子代tgid的行为只针对helloTree
	// q不添加，直接把新的tgid对应的tgidNode入队就是了
	q.Enqueue(tmp)
	visited[starTgid] = true
	for !q.IsEmpty() {
		tmp, ok := q.Dequeue()
		if !ok {
			continue
		}
		node := tmp.(*tgidNode) // 队列里的一个节点，这里必须重新申请node
		helloTree[node.Tgid] = node
		for i := 0; i < len(node.Threads); i++ {
			for j := 0; j < len(node.Threads[i].Children); j++ {
				tgid := findTgid[node.Threads[i].Children[j]]
				_, exists := visited[tgid]
				if !exists {
					// 子代里有没见过的tgid
					tgidNode, exists := tgidMap[tgid]
					if !exists {
						continue
					}
					helloTree[node.Tgid].ChildTgid = append(helloTree[node.Tgid].ChildTgid, tgid)
					q.Enqueue(tgidNode)
					visited[tgid] = true
				}
			}
		}
	}
}

func optimazePid(starTgid int, rootfsPids []int) {
	getDockerRootFs := make(map[string]string) // dockerId --> rootfs
	// 首先处理一下记录有pivot_root信息的进程，防止pivot先于fork
	for _, rootfsPid := range rootfsPids {
		rootfsTgid := findTgid[rootfsPid]
		i := helloTree[rootfsTgid].FindPid[rootfsPid]
		rootfsProcess := &(helloTree[rootfsTgid].Threads[i])
		if rootfsProcess.RootFS == "cwd" {
			rootfsProcess.RootFS = rootfsProcess.Cwd
		}
		getDockerRootFs[rootfsProcess.DockerId] = rootfsProcess.RootFS
	}

	count = 0
	for _, val := range helloTree {
		// 处理一下pid结束时间，顺便找找爹
		// 结束时间是因为很多线程结束时间没获取到，默认按照进程退出时间处理
		// Ppid是因为进程产生之初收到的信息写的爹一定是亲爹
		// 但是产生线程时候该进程很可能已作为孤儿被收养，导致线程里关于爹的记录是继父
		for i := 0; i < len(val.Threads); i++ {
			if i != 0 {
				if val.Threads[i].Tgid < val.Threads[0].Tgid {
					val.Threads[i].ParentTgid = val.Threads[0].ParentTgid
					val.Threads[i].Ppid = val.Threads[0].Ppid
				}
				if val.Threads[i].ExitTimestamp.IsZero() {
					val.Threads[i].ExitCode = val.Threads[0].ExitCode
					val.Threads[i].ExitTimestamp = val.Threads[0].ExitTimestamp
					val.Threads[i].ExitSignal = val.Threads[0].ExitSignal
				}
			}

			dockerId := val.Threads[i].DockerId
			if dockerId != "" {
				rootfs, exists := getDockerRootFs[dockerId]
				if !exists {
					fmt.Fprintf(os.Stderr, "Err: the docker rootfs of pid %d is not known!\n", val.Threads[i].Pid)
					continue
				}
				val.Threads[i].RootFS = rootfs
			}
		}

		count++
	}
}

// 绘制进程树
func drawTree(treeFile *os.File, pidFile *os.File, node *tgidNode, prefix string, isLast bool) {
	if node == nil {
		return
	}

	fmt.Fprintf(treeFile, "%s", prefix)
	if isLast {
		fmt.Fprintf(treeFile, "└── ")
		prefix += "    "
	} else {
		fmt.Fprintf(treeFile, "├── ")
		prefix += "│   "
	}
	// 将当前进程的参数整理为一行命令
	argv := ""
	for i, arg := range node.Threads[0].Args {
		if i == 0 {
			argv = arg
		} else {
			argv += " " + arg
		}
	}
	fmt.Fprintf(treeFile, "%d: %s\n", node.Tgid, argv)

	// 当前节点信息以json格式写入pidFile
	jsonData, err := json.MarshalIndent(node, "", "    ")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Err: %v\n", err)
		return
	}
	pidFile.Write(jsonData)
	pidFile.WriteString("\n\n")

	// 递归打印子节点
	for i, childTgid := range node.ChildTgid {
		childNode, exists := helloTree[childTgid]
		if exists {
			drawTree(treeFile, pidFile, childNode, prefix, i == len(node.ChildTgid)-1)
		}
	}
}