package cgroup

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
)

// hostProcsMaxDepth bounds the walk below a host's cgroup. The platform
// nests at most a few levels (exec group, role child, generation); a tree
// deeper than this is not one it made, and is reported rather than walked.
const hostProcsMaxDepth = 16

func hostProcs(root, host, unit string) (int, error) {
	if root == "" {
		return 0, ErrUnknown
	}
	if host == "" || strings.ContainsAny(host, "/\x00") || host == "." || host == ".." {
		return 0, fmt.Errorf("cgroup: %q is not a host name", host)
	}
	n, err := countProcs(filepath.Join(root, ExecName(host)), 0)
	if err != nil {
		return 0, err
	}
	if unit != "" {
		if !filepath.IsAbs(unit) {
			return 0, fmt.Errorf("cgroup: the unit cgroup %q is not an absolute path", unit)
		}
		m, err := countProcs(filepath.Clean(unit), 0)
		if err != nil {
			return 0, err
		}
		n += m
	}
	return n, nil
}

// countProcs is the number of pids listed in dir's cgroup.procs and in every
// cgroup below it. An absent dir is 0. Symlinks are never followed: cgroupfs
// has none, so a link is not a cgroup.
func countProcs(dir string, depth int) (int, error) {
	if depth > hostProcsMaxDepth {
		return 0, fmt.Errorf("cgroup: %s is nested deeper than %d levels", dir, hostProcsMaxDepth)
	}
	st, err := os.Lstat(dir)
	if errors.Is(err, os.ErrNotExist) {
		return 0, nil
	}
	if err != nil {
		return 0, err
	}
	if !st.IsDir() {
		return 0, fmt.Errorf("cgroup: %s is not a cgroup directory", dir)
	}
	n, err := readProcsFile(filepath.Join(dir, "cgroup.procs"))
	if err != nil {
		return 0, err
	}
	ents, err := os.ReadDir(dir)
	if errors.Is(err, os.ErrNotExist) {
		return n, nil // removed while read: its processes are gone with it
	}
	if err != nil {
		return 0, err
	}
	for _, e := range ents {
		if !e.IsDir() || e.Type()&os.ModeSymlink != 0 {
			continue
		}
		m, err := countProcs(filepath.Join(dir, e.Name()), depth+1)
		if err != nil {
			return 0, err
		}
		n += m
	}
	return n, nil
}

// readProcsFile counts the non-empty lines of a cgroup.procs file. A file
// that is gone (the cgroup was removed between the listing and the read)
// lists nothing.
func readProcsFile(path string) (int, error) {
	f, err := os.Open(path)
	if errors.Is(err, os.ErrNotExist) {
		return 0, nil
	}
	if err != nil {
		return 0, err
	}
	defer f.Close()
	n := 0
	sc := bufio.NewScanner(f)
	for sc.Scan() {
		if strings.TrimSpace(sc.Text()) != "" {
			n++
		}
	}
	if err := sc.Err(); err != nil {
		// A cgroup removed while it was read answers ENODEV.
		if _, statErr := os.Stat(filepath.Dir(path)); errors.Is(statErr, os.ErrNotExist) {
			return 0, nil
		}
		return 0, err
	}
	return n, nil
}