package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"sync"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

const (
	oldDBName     = "test"
	oldPidColName = "pids"
)

type treeNode struct {
	Tgid     int
	Threads  []Process
	Children []int
}

func main() {
	// 连接到MongoDB
	client, err := mongo.Connect(context.TODO(), options.Client().ApplyURI("mongodb://localhost:27017"))
	if err != nil {
		log.Fatal(err)
	}
	defer client.Disconnect(context.TODO())

	// 选择数据库和集合
	db := client.Database(oldDBName)
	collection := db.Collection(oldPidColName)

	// 提取所有数据
	var res []Process

	cursor, err := collection.Find(context.Background(), bson.M{})
	if err != nil {
		fmt.Fprintf(os.Stderr, "Err: %v\n", err)
	}
	if err := cursor.All(context.Background(), &res); err != nil {
		fmt.Fprintf(os.Stderr, "Err All: %v\n", err)
	}

	var merged sync.Map
	for _, process := range res {
		tmp, ok := merged.Load(process.Pid)
		if ok {
			// 证明重复了，要合并
			tmp := ProMerge(tmp.(Process), process)
			merged.Store(process.Pid, tmp)
		} else {
			// 没有，直接插入
			merged.Store(process.Pid, process)
		}
	}

	var treeMap sync.Map
	findTgid := make(map[int]int)
	var stared int
	merged.Range(func(key, val interface{}) bool {
		tmp := val.(Process)
		if tmp.Star {
			stared = tmp.Tgid
		}
		// 登记tgid
		findTgid[tmp.Pid] = tmp.Tgid
		nodeTmp, ok := treeMap.Load(tmp.Tgid)
		if ok {
			// 直接记录
			node := nodeTmp.(treeNode)
			node.Threads = append(node.Threads, tmp)
			node.Children = append(node.Children, tmp.Children...)
			treeMap.Store(tmp.Tgid, node)
		} else {
			node := treeNode{
				Tgid:     tmp.Tgid,
				Threads:  make([]Process, 0),
				Children: make([]int, 0),
			}
			node.Threads = append(node.Threads, tmp)
			node.Children = append(node.Children, tmp.Children...)
			treeMap.Store(tmp.Tgid, node)
		}
		return true
	})

	// 从tgid==stared开始，构建树
	var helloTree sync.Map // 在树上的tgid节点
	var q Queue            // 记录每一个整理好的结构体，bfs
	visited := make(map[int]bool)
	visited[stared] = true
	tmp, ok := treeMap.Load(stared)
	if !ok {
		return
	}

	q.Enqueue(tmp)
	helloTree.Store(stared, tmp)
	for !q.IsEmpty() {
		tmp, ok := q.Dequeue()
		if !ok {
			continue
		}
		node := tmp.(treeNode)
		for i := 0; i < len(node.Children); i++ {
			tgid := findTgid[node.Children[i]]
			_, exists := visited[tgid]
			if !exists {
				visited[tgid] = true
				tgidNode, ok := treeMap.Load(tgid)
				if !ok {
					continue
				}
				helloTree.Store(tgid, tgidNode)
				q.Enqueue(tgidNode)
			}
		}
	}

	// TODO:
	// 1.修改数据结构，使之自身即存储树结构，插入数据库后前端拿出来就能用
	// 2.还有其余优化要做，比如线程退出时间与进程推出时间

	// count := 0
	// helloTree.Range(func(key, val interface{}) bool {
	// 	count++
	// 	fmt.Printf("tgid: %d\n", val.(treeNode).Tgid)
	// 	return true
	// })
	// fmt.Printf("Star: %d, res: %d\n", stared, count)

	// 接下来处理文件
}

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
}