package procgroup

import (
	"context"
	"os/exec"
	"strconv"
	"strings"
	"syscall"
)

type Cmd struct {
	*exec.Cmd
}

func Command(name string, args ...string) *Cmd {
	return group(exec.Command(name, args...))
}

func CommandContext(ctx context.Context, name string, args ...string) *Cmd {
	cmd := group(exec.CommandContext(ctx, name, args...))
	cmd.Cancel = cmd.Kill
	return cmd
}

func group(cmd *exec.Cmd) *Cmd {
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	return &Cmd{Cmd: cmd}
}

// Kill also kills descendants that left the group: agents run their commands in sessions of their own.
func (cmd *Cmd) Kill() error {
	pids := descendants(cmd.Process.Pid)
	err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
	for _, pid := range pids {
		syscall.Kill(pid, syscall.SIGKILL)
	}
	return err
}

func descendants(root int) []int {
	out, err := exec.Command("ps", "-A", "-o", "pid=,ppid=").Output()
	if err != nil {
		return nil
	}
	children := map[int][]int{}
	for _, line := range strings.Split(string(out), "\n") {
		f := strings.Fields(line)
		if len(f) != 2 {
			continue
		}
		pid, err1 := strconv.Atoi(f[0])
		ppid, err2 := strconv.Atoi(f[1])
		if err1 != nil || err2 != nil {
			continue
		}
		children[ppid] = append(children[ppid], pid)
	}
	pids := []int{}
	queue := children[root]
	for len(queue) > 0 {
		pid := queue[0]
		queue = queue[1:]
		pids = append(pids, pid)
		queue = append(queue, children[pid]...)
	}
	return pids
}