//go:build linux

package cgroup

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"slices"
	"strconv"
	"strings"
	"time"

	"github.com/carlosframework/platform/internal/controlname"
	"golang.org/x/sys/unix"
)

type Cleanup struct{ state *cleanupState }
type cleanupState struct {
	pin   *Prepared
	procs *os.File
	// m and path are the manager that opened the cleanup and the directory
	// it pinned, as the manager's starts name it: the in-flight registry's
	// key (inflight.go). A process being placed there or below counts as a
	// member (Ruling 6).
	m    *Manager
	path string
}

func cleanupFilesystem(root *os.File) error {
	var st unix.Statfs_t
	if err := unix.Fstatfs(int(root.Fd()), &st); err != nil {
		return err
	}
	if uint64(st.Type) != uint64(unix.CGROUP2_SUPER_MAGIC) {
		return ErrStrictUnsupported
	}
	return nil
}

func (m *Manager) OpenCleanup(host string) (*Cleanup, error) {
	return m.openCleanupAt(host, "", cleanupFilesystem)
}

// OpenCleanupChild pins one child of host's generation parent (zero-
// downtime spec §7.3): its KillAndEmpty writes that child's cgroup.kill
// and proves that child empty, and never touches a sibling. ErrNoCgroup
// when the child is not there.
func (m *Manager) OpenCleanupChild(host, child string) (*Cleanup, error) {
	return m.openCleanupAt(host, child, cleanupFilesystem)
}

func (m *Manager) openCleanup(host string, checkFS func(*os.File) error) (*Cleanup, error) {
	return m.openCleanupAt(host, "", checkFS)
}

func (m *Manager) openCleanupAt(host, child string, checkFS func(*os.File) error) (out *Cleanup, err error) {
	if m == nil || m.root == "" {
		return nil, ErrNoDelegation
	}
	if !controlname.ValidHostShape(host) || (child != "" && !isChildName(child)) {
		return nil, ErrStrictRefused
	}
	root, err := os.OpenFile(m.root, os.O_RDONLY|unix.O_DIRECTORY|unix.O_NOFOLLOW, 0)
	if err != nil {
		return nil, err
	}
	p := &Prepared{strictPreparedState: &strictPreparedState{root: root, host: host}}
	path := filepath.Join(m.root, ExecName(host))
	if child != "" {
		path = filepath.Join(path, child)
	}
	c := &Cleanup{state: &cleanupState{pin: p, m: m, path: path}}
	defer func() {
		if err != nil {
			if c.state.procs != nil {
				c.state.procs.Close()
			}
			p.closeFiles()
		}
	}()
	if err = checkFS(root); err != nil {
		return nil, err
	}
	p.dir, err = strictOpen(root, ExecName(host), unix.O_RDONLY|unix.O_DIRECTORY)
	if errors.Is(err, unix.ENOENT) {
		return nil, ErrNoCgroup
	}
	if err != nil {
		return nil, err
	}
	if child != "" {
		// The child is pinned beneath the parent's own descriptor, and its
		// identity is checked against that descriptor from now on.
		p.entryParent, p.entryName = p.dir, child
		p.dir, err = strictOpen(p.entryParent, child, unix.O_RDONLY|unix.O_DIRECTORY)
		if errors.Is(err, unix.ENOENT) {
			return nil, ErrNoCgroup
		}
		if err != nil {
			return nil, err
		}
	}
	p.identity, err = strictID(p.dir)
	if err != nil {
		return nil, err
	}
	p.events, err = strictOpen(p.dir, "cgroup.events", unix.O_RDONLY)
	if err != nil {
		return nil, err
	}
	p.kill, err = strictOpen(p.dir, "cgroup.kill", unix.O_WRONLY)
	if errors.Is(err, unix.ENOENT) {
		// The cgroup is there (its events were opened above), so a missing
		// cgroup.kill is the kernel's, not a cgroup that went away.
		return nil, fmt.Errorf("%w: %w", ErrNoKill, err)
	}
	if err != nil {
		return nil, err
	}
	c.state.procs, err = strictOpen(p.dir, "cgroup.procs", unix.O_RDONLY)
	if err != nil {
		return nil, err
	}
	if err = p.identityOK(); err != nil {
		return nil, err
	}
	return c, nil
}

func (c *Cleanup) Populated() (bool, error) {
	if c == nil || c.state == nil {
		return false, ErrStrictState
	}
	p := c.state.pin
	p.mu.Lock()
	defer p.mu.Unlock()
	return c.populationLocked()
}

// populationLocked is the pinned cgroup's population. A process being
// placed there or below it (the in-flight registry) makes it populated: a
// shim spends its first milliseconds in the agent's cgroup, and a flat
// fallback's child runs there until its PID is moved, so the kernel cannot
// see either yet. cgroup.events is read inside readQuiet, so a placement
// that registers during the read, or lands and ends its registration
// between the read and the registry's, still counts (whole-branch review,
// finding 1).
func (c *Cleanup) populationLocked() (bool, error) {
	p := c.state.pin
	if err := p.identityOK(); err != nil {
		return false, err
	}
	var (
		v   bool
		err error
	)
	quiet := c.state.m.readQuiet(c.state.path, func() { v, err = strictPopulation(p.events) })
	if err != nil {
		return false, err
	}
	if err := p.identityOK(); err != nil {
		return false, err
	}
	return v || !quiet, nil
}

func (c *Cleanup) Members() ([]int, error) {
	if c == nil || c.state == nil {
		return nil, ErrStrictState
	}
	p := c.state.pin
	p.mu.Lock()
	defer p.mu.Unlock()
	return c.membersLocked()
}

// Cleanup's walk is bounded: a tree deeper or wider than this is refused
// rather than vouched for unread.
const (
	maxCleanupDepth  = 8
	maxCleanupGroups = 1024
)

// membersLocked lists every process in the pinned cgroup AND in every
// cgroup below it. A generation parent holds no process itself (zero-
// downtime spec §7.3), so stop evidence taken on it must read its
// generations'. A flat tenant runs as the agent's uid and can make cgroups
// of its own under its generation, so the walk goes to any depth, bounded.
// While a process is being placed there or below it, its PID is not in any
// cgroup.procs yet, so the list would be short: ErrPlacementInFlight. The
// walk runs inside readQuiet, so a placement in flight at any moment of it
// counts, not only one in flight before it began.
func (c *Cleanup) membersLocked() ([]int, error) {
	p := c.state.pin
	if err := p.identityOK(); err != nil {
		return nil, err
	}
	var (
		members []int
		err     error
	)
	quiet := c.state.m.readQuiet(c.state.path, func() {
		if members, err = readProcs(c.state.procs); err != nil {
			return
		}
		groups := 0
		err = walkMembers(p.dir, 1, &groups, &members)
	})
	if err != nil {
		return nil, err
	}
	if !quiet {
		return nil, ErrPlacementInFlight
	}
	slices.Sort(members)
	if err := p.identityOK(); err != nil {
		return nil, err
	}
	return members, nil
}

// readProcs reads every PID in one cgroup.procs. It may exceed
// strictRead's 4096-byte bound, so every line is scanned.
func readProcs(f *os.File) ([]int, error) {
	if _, err := f.Seek(0, 0); err != nil {
		return nil, err
	}
	scan := bufio.NewScanner(f)
	var pids []int
	for scan.Scan() {
		raw := strings.TrimSpace(scan.Text())
		pid, err := strconv.Atoi(raw)
		if err != nil || pid <= 0 {
			return nil, fmt.Errorf("%w: invalid cgroup member", ErrStrictRefused)
		}
		pids = append(pids, pid)
	}
	return pids, scan.Err()
}

// cleanupReadBatch is how many directory entries one read takes, so the
// group budget is checked as the listing is read, not after all of it.
const cleanupReadBatch = 64

// walkMembers appends the members of every cgroup below dir. A child that
// vanishes between the listing and its open was removed, so it is empty.
func walkMembers(dir *os.File, depth int, groups *int, members *[]int) error {
	names, err := childGroups(dir, depth, groups)
	if err != nil {
		return err
	}
	for _, name := range names {
		child, err := strictOpen(dir, name, unix.O_RDONLY|unix.O_DIRECTORY)
		if errors.Is(err, unix.ENOENT) {
			continue
		}
		if err != nil {
			return err
		}
		procs, err := strictOpen(child, "cgroup.procs", unix.O_RDONLY)
		if errors.Is(err, unix.ENOENT) {
			child.Close()
			continue
		}
		if err != nil {
			child.Close()
			return err
		}
		pids, err := readProcs(procs)
		procs.Close()
		if err == nil {
			*members = append(*members, pids...)
			err = walkMembers(child, depth+1, groups, members)
		}
		child.Close()
		if err != nil {
			return err
		}
	}
	return nil
}

// childGroups lists the child cgroups of dir in bounded batches, counting
// each against the walk's budget as it is read. A tree past the depth or
// group bound is refused before any more of it is read, so a tenant's many
// empty cgroups cost at most the budget.
func childGroups(dir *os.File, depth int, groups *int) ([]string, error) {
	fd, err := unix.Openat(int(dir.Fd()), ".", unix.O_RDONLY|unix.O_DIRECTORY|unix.O_CLOEXEC, 0)
	if err != nil {
		return nil, err
	}
	list := os.NewFile(uintptr(fd), dir.Name())
	defer list.Close()
	var names []string
	for {
		entries, err := list.ReadDir(cleanupReadBatch)
		for _, e := range entries {
			if !e.IsDir() {
				continue
			}
			if depth > maxCleanupDepth {
				return nil, fmt.Errorf("%w: cgroup tree deeper than %d", ErrStrictRefused, maxCleanupDepth)
			}
			if *groups++; *groups > maxCleanupGroups {
				return nil, fmt.Errorf("%w: more than %d cgroups below one host", ErrStrictRefused, maxCleanupGroups)
			}
			names = append(names, e.Name())
		}
		if errors.Is(err, io.EOF) {
			return names, nil
		}
		if err != nil {
			return nil, err
		}
	}
}

// cleanupKillWrite writes a pinned cgroup's cgroup.kill: a seam so a test
// can act right after the kill lands.
var cleanupKillWrite = func(f *os.File) error { return strictWrite(f, "1") }

// KillAndEmpty kills everything in the pinned cgroup and below it and
// returns nil only once it is proven empty: unpopulated and no members.
// Any error is unproven, and ctx must carry a deadline.
//
// A process being placed into the cgroup or below it (the in-flight
// registry, Ruling 6) is a member the kernel cannot see yet. KillAndEmpty
// first waits for every such placement to resolve, bounded by ctx; an
// error there is unproven (ErrPlacementInFlight with ctx's error), and
// nothing is killed. cgroup.kill is then written once. Every emptiness
// check after it reads the kernel inside readQuiet, so a placement in
// flight at any moment of a check makes that check populated. A placement
// that registers after the wait is therefore never missed, but it is not
// killed either: the kill was written before it landed, and is not written
// again. Its process keeps the cgroup populated, so the proof ends
// unproven when ctx ends, unless that process exits first. Never a false
// empty; a late placement costs the proof, not the tenant. Callers still
// exclude new starts during a proof as they always have (the lease, and
// the before hook a start asks immediately before exec); the registry only
// closes the window a shim, or a flat start awaiting its PID move, spends
// outside its cgroup.
func (c *Cleanup) KillAndEmpty(ctx context.Context) error {
	if c == nil || c.state == nil {
		return ErrStrictState
	}
	if ctx == nil {
		return ErrStrictDeadline
	}
	deadline, ok := ctx.Deadline()
	if !ok {
		return ErrStrictDeadline
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	if !time.Now().Before(deadline) {
		return context.DeadlineExceeded
	}
	if err := c.state.m.waitPlacementsUnder(ctx, c.state.path); err != nil {
		return fmt.Errorf("%w: %w", ErrPlacementInFlight, err)
	}
	p := c.state.pin
	lockTick := time.NewTicker(time.Millisecond)
	defer lockTick.Stop()
	for !p.mu.TryLock() {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-lockTick.C:
		}
	}
	defer p.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return err
	}
	if err := p.identityOK(); err != nil {
		return err
	}
	// Marked first: the next start into this cgroup makes it fresh, since
	// the kernel kills anything cloned into one that took cgroup.kill
	// (amadan #272, recreate_linux.go).
	markKilled(p.identity)
	if err := cleanupKillWrite(p.kill); err != nil {
		return err
	}
	tick := time.NewTicker(5 * time.Millisecond)
	defer tick.Stop()
	for {
		if err := ctx.Err(); err != nil {
			return err
		}
		populated, err := c.populationLocked()
		if err != nil {
			return err
		}
		members, err := c.membersLocked()
		if errors.Is(err, ErrPlacementInFlight) {
			// In flight during the members read: not empty yet.
			populated, members, err = true, nil, nil
		}
		if err != nil {
			return err
		}
		if !populated && len(members) == 0 {
			return ctx.Err()
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-tick.C:
		}
	}
}

func (c *Cleanup) Close() error {
	if c == nil || c.state == nil {
		return ErrStrictState
	}
	p := c.state.pin
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.closed {
		return nil
	}
	p.closed = true
	var errs []error
	for _, f := range []*os.File{c.state.procs, p.events, p.kill, p.dir, p.entryParent, p.root} {
		if f != nil {
			errs = append(errs, f.Close())
		}
	}
	return errors.Join(errs...)
}