package main

import (
	"fmt"
	"syscall"

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

const (
	dbName      string = "test"
	pidColName  string = "pids"
	fdColName   string = "fds"
	fileColName string = "files"
)

// var mongoMutex sync.Mutex
var pidCol, fdCol, fileCol mongoClient

var docRes []bson.M
var err error

func deal() {
	defer wg.Done()
	var cooked Event
	var ok bool

	if err = pidCol.init(dbName, pidColName); err != nil {
		fmt.Printf("Error while initing the mongodb: %v\n", err)
		return
	}
	err = pidCol.InsertOne(bson.M{
		"ppid":     1,
		"pid":      containerdPid,
		"cwd":      "/",
		"children": []bson.M{},
	})
	if err != nil {
		fmt.Printf("Error while initing the mongodb: %v\n", err)
		return
	}

	if err = fdCol.init(dbName, fdColName); err != nil {
		fmt.Printf("Error while initing the mongodb: %v\n", err)
		return
	}
	if err = fileCol.init(dbName, fileColName); err != nil {
		fmt.Printf("Error while initing the mongodb: %v\n", err)
	}

	fmt.Printf("Containerd: %d\n", containerdPid)
	defer pidCol.Disconnect()
	defer fdCol.Disconnect()
	defer fileCol.Disconnect()

	for {
		cooked, ok = <-cookedChan
		if !ok {
			break
		}

		switch cooked.tag {
		case NEWPID:
			go dealNewPid(cooked)
		case EXECVE:
			go dealExecve(cooked)
		case PIDEXIT:
			go deletePid(cooked)
		case FILEOPEN:
			fileOpen(cooked)
		case FILEWRITE:
			fileWrite(cooked)
		case FILECLOSE:
			fileClose(cooked)
		}
	}
}

func deletePid(cooked Event) {
	pidCol.GetLock()
	// 先从老爹那里销户
	pidCol.UpdateOne(bson.M{"pid": cooked.ppid}, bson.M{
		"$pull": bson.M{
			"children": cooked.pid,
		},
	})

	// 孩子们需要收容
	// 不必到children里一个个找，直接看ppid即可
	// pidCol.UpdateMany(bson.M{"ppid": cooked.pid}, bson.M{"ppid": 1})
	// 在这套逻辑里，孩子是不需要收容的，因为我们根本就不看ppid来工作

	// 可以去死了
	pidCol.UpdateOne(bson.M{"pid": cooked.pid}, bson.M{
		"$set": bson.M{
			"exit_timestamp": cooked.timestamp,
			"exit_code":      cooked.exit_code,
			"exit_signal":    cooked.exit_signal,
		},
	})
	pidCol.Mutex.Unlock()
}

func dealNewPid(cooked Event) {
	fmt.Printf("Fork: %v\t%6d\t%6d\n", cooked.timestamp, cooked.ppid, cooked.pid)
	// 有无父进程在观察中
	docRes, err = pidCol.Finddoc(bson.M{"pid": cooked.ppid})
	if err != nil || len(docRes) != 1 {
		return
	}

	// 自身是否已经记录
	docRes, err = pidCol.Finddoc(bson.M{"pid": cooked.pid})
	if err != nil {
		fmt.Printf("Err finding: %v\n", err)
		return
	}

	pidCol.GetLock()
	if len(docRes) != 0 {
		// 进程原本就存在，换言之别的消息先到了
		// 所有先行抵达的消息必须保留execve/children字段
		// 此处不再更新
		// 以防把原有信息更没了
		pidCol.UpdateOne(bson.M{"pid": cooked.pid}, bson.M{
			"start_timestamp": cooked.timestamp,
			"ppid":            cooked.ppid,
			"pid":             cooked.pid,
			"cwd":             cooked.cwd,
			// "execve":          []bson.M{},
			"args": cooked.argv,
			// "children":        []bson.M{},
		})
	} else {
		// 这进程本是新修的
		pidCol.InsertOne(bson.M{
			"start_timestamp": cooked.timestamp,
			"ppid":            cooked.ppid,
			"pid":             cooked.pid,
			"cwd":             cooked.cwd,
			"execve":          []bson.M{},
			"args":            cooked.argv,
			"children":        []bson.M{},
		})
	}

	pidCol.UpdateOne(bson.M{"pid": cooked.ppid}, bson.M{
		"$push": bson.M{
			"children": cooked.pid,
		},
	})
	pidCol.Mutex.Unlock()
}

func dealExecve(cooked Event) {
	fmt.Printf("EXEC: %6d\t%6d\n", cooked.ppid, cooked.pid)
	// 父进程在不在？不在扔
	docRes, err = pidCol.Finddoc(bson.M{"pid": cooked.ppid})
	if err != nil || len(docRes) != 1 {
		return
	}

	// 首先检查进程是否存在，如不存在则为之创建
	docRes, err = pidCol.Finddoc(bson.M{"pid": cooked.pid})
	if err != nil {
		return
	}

	pidCol.GetLock()
	if len(docRes) == 1 {
		// 自身已在，直接记录
		pidCol.UpdateOne(bson.M{"pid": cooked.pid}, bson.M{
			"$push": bson.M{
				"execve": bson.M{
					"timestamp": cooked.timestamp,
					"execArgs":  cooked.argv,
				},
			},
		})
	} else {
		// 先fork抵达，插入
		pidCol.InsertOne(bson.M{
			"ppid":     cooked.ppid,
			"pid":      cooked.pid,
			"children": []bson.M{},
			"execve": []bson.M{
				{
					"timestamp": cooked.timestamp,
					"execArgs":  cooked.argv,
				},
			},
		})
	}
	pidCol.Mutex.Unlock()
}

func fileOpen(cooked Event) {
	// 权限检查过了，不必再查
	fdCol.InsertOne(bson.M{
		"timestamp": cooked.timestamp,
		"fileName":  cooked.pathName,
		"pid":       cooked.pid,
		"fd":        cooked.exit_code,
		"flags":     cooked.syscallParam,
		"written":   []bson.M{},
	})

	if cooked.syscallParam[1]&syscall.O_TRUNC != 0 {
		fdCol.UpdateOne(bson.M{"pid": cooked.pid, "fd": cooked.exit_code}, bson.M{
			"$push": bson.M{
				"written": cooked.timestamp,
			},
		})
	}
}

func fileClose(cooked Event) {
	res, err := fdCol.Finddoc(bson.M{
		"pid":             cooked.pid,
		"fd":              cooked.syscallParam[0],
		"close_timestamp": bson.M{"$exists": false},
	})
	if err != nil {
		fmt.Printf("Err closing fd %d of pid %d: %v\n", cooked.syscallParam[0], cooked.pid, err)
	}
	if len(res) == 0 {
		return
	}
	fdCol.UpdateOne(bson.M{
		"pid":             cooked.pid,
		"fd":              cooked.syscallParam[0],
		"close_timestamp": bson.M{"$exists": false},
	}, bson.M{"$set": bson.M{"close_timestamp": cooked.timestamp}})
}

func fileWrite(cooked Event) {
	res, err := fdCol.Finddoc(bson.M{
		"pid":             cooked.pid,
		"fd":              cooked.syscallParam[0],
		"close_timestamp": bson.M{"$exists": false},
	})
	if err != nil {
		fmt.Printf("Err closing fd %d of pid %d: %v\n", cooked.syscallParam[0], cooked.pid, err)
	}
	if len(res) == 0 {
		return
	}
	fdCol.UpdateOne(bson.M{
		"pid":             cooked.pid,
		"fd":              cooked.syscallParam[0],
		"close_timestamp": bson.M{"$exists": false},
	}, bson.M{"$push": bson.M{"written": cooked.timestamp}})
}