package main

import (
	"unsafe"

	"golang.org/x/sys/windows"
)

type (
	ProcessMap struct {
		Map       map[uint32]*ProcessMapNode
		NameIndex map[string][]*ProcessMapNode
	}
	ProcessMapNode struct {
		Proc     windows.ProcessEntry32
		Name     string
		Children map[uint32]*ProcessMapNode
	}
)

func (pt *ProcessMap) add(proc *windows.ProcessEntry32) error {
	if pt.Map == nil {
		pt.Map = make(map[uint32]*ProcessMapNode)
	}
	if pt.NameIndex == nil {
		pt.NameIndex = make(map[string][]*ProcessMapNode)
	}

	procNode := &ProcessMapNode{
		Proc: *proc,
		Name: windows.UTF16ToString(proc.ExeFile[:]),
	}

	_, ok := pt.NameIndex[procNode.Name]
	if !ok {
		pt.NameIndex[procNode.Name] = make([]*ProcessMapNode, 0)
	}
	pt.NameIndex[procNode.Name] = append(pt.NameIndex[procNode.Name], procNode)

	parent, parentExists := pt.Map[proc.ParentProcessID]
	if !parentExists {
		parent = procNode
		pt.Map[proc.ParentProcessID] = parent
	}

	child, childExists := pt.Map[proc.ProcessID]
	if !childExists {
		child = procNode
		pt.Map[proc.ProcessID] = child
	}

	if parent.Children == nil {
		parent.Children = make(map[uint32]*ProcessMapNode)
	}
	parent.Children[proc.ProcessID] = child

	return nil
}

func (pt *ProcessMap) findByName(name string) ([]*ProcessMapNode, bool) {
	val, ok := pt.NameIndex[name]
	return val, ok
}
func (pt *ProcessMap) findByPid(pid uint32) (*ProcessMapNode, bool) {
	val, ok := pt.Map[pid]
	return val, ok
}

func BuildProcessMap() (*ProcessMap, error) {
	tree := &ProcessMap{}

	snapshot, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPPROCESS, 0)
	if err != nil {
		return nil, err
	}
	defer windows.CloseHandle(snapshot)

	var pe32 windows.ProcessEntry32
	pe32.Size = uint32(unsafe.Sizeof(pe32))

	var i int
	err = windows.Process32First(snapshot, &pe32)
	for err == nil {
		tree.add(&pe32)
		i++
		if i > 500 {
			break
		}
		err = windows.Process32Next(snapshot, &pe32)
	}
	return tree, nil
}