//go:build linux

package cgroup

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

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

type strictIdentity struct{ device, inode uint64 }

var strictReservations = struct {
	sync.Mutex
	held map[strictIdentity]bool
}{held: make(map[strictIdentity]bool)}

// Prepared pins one empty host cgroup. It grants no lease, fence or database
// authority. The caller must also obtain Cmd.Wait evidence after cleanup.
type Prepared struct{ *strictPreparedState }

type strictPreparedState struct {
	mu                       sync.Mutex
	root, dir, kill, events  *os.File
	limits                   map[string]*os.File
	expected                 map[string]string
	identity                 strictIdentity
	entryParent              *os.File
	entryName                string
	parent                   *Prepared
	draining                 atomic.Bool
	host                     string
	attempted, empty, closed bool
	start                    func(*exec.Cmd) error
	// genParent marks a generation parent (PrepareStrictParent): it holds
	// no process, pins PrepareGen children, and moves its limit only
	// through SetParentLimit. tier is the tier it was prepared with, or a
	// generation child's own.
	genParent bool
	tier      Limit
}

func strictOpen(dir *os.File, name string, flags int) (*os.File, error) {
	fd, e := unix.Openat(int(dir.Fd()), name, flags|unix.O_CLOEXEC|unix.O_NOFOLLOW, 0)
	if e != nil {
		return nil, e
	}
	return os.NewFile(uintptr(fd), name), nil
}
func strictID(f *os.File) (strictIdentity, error) {
	var st unix.Stat_t
	if e := unix.Fstat(int(f.Fd()), &st); e != nil {
		return strictIdentity{}, e
	}
	return strictIdentity{uint64(st.Dev), st.Ino}, nil
}
func strictRead(f *os.File) (string, error) {
	if _, e := f.Seek(0, 0); e != nil {
		return "", e
	}
	b, e := io.ReadAll(io.LimitReader(f, 4097))
	if e != nil || len(b) > 4096 {
		return "", ErrStrictRefused
	}
	return strings.TrimSpace(string(b)), nil
}
func strictWrite(f *os.File, value string) error {
	if _, e := f.Seek(0, 0); e != nil {
		return e
	}
	n, e := f.WriteString(value)
	if e != nil {
		return e
	}
	if n != len(value) {
		return io.ErrShortWrite
	}
	return nil
}
func strictPopulation(f *os.File) (bool, error) {
	raw, e := strictRead(f)
	if e != nil {
		return false, e
	}
	seen := false
	populated := false
	for _, line := range strings.Split(raw, "\n") {
		parts := strings.Fields(line)
		if len(parts) != 2 {
			return false, ErrStrictRefused
		}
		if parts[0] == "populated" {
			if seen || (parts[1] != "0" && parts[1] != "1") {
				return false, ErrStrictRefused
			}
			seen = true
			populated = parts[1] == "1"
		}
	}
	if !seen {
		return false, ErrStrictRefused
	}
	return populated, nil
}
func (p *Prepared) identityOK() error {
	if p.closed {
		return ErrStrictState
	}
	id, e := strictID(p.dir)
	if e != nil || id != p.identity {
		return ErrStrictRefused
	}
	var st unix.Stat_t
	parent, name := p.root, ExecName(p.host)
	if p.entryParent != nil {
		parent, name = p.entryParent, p.entryName
	}
	if e = unix.Fstatat(int(parent.Fd()), name, &st, unix.AT_SYMLINK_NOFOLLOW); e != nil || st.Mode&unix.S_IFMT != unix.S_IFDIR || (strictIdentity{uint64(st.Dev), st.Ino}) != p.identity {
		return ErrStrictRefused
	}
	return nil
}
func (p *Prepared) limitsOK() error {
	for name, want := range p.expected {
		got, e := strictRead(p.limits[name])
		if e != nil || got != want {
			return ErrStrictRefused
		}
	}
	return nil
}
func (p *Prepared) closeFiles() {
	for _, f := range p.limits {
		f.Close()
	}
	for _, f := range []*os.File{p.events, p.kill, p.dir, p.entryParent, p.root} {
		if f != nil {
			f.Close()
		}
	}
}

// selfCgroupFile is where the strict preflight reads this process's own
// cgroup. A variable so a test can place the agent where it needs.
var selfCgroupFile = "/proc/self/cgroup"

// strictRoot is the host-independent half of PrepareStrict's runtime
// checks, and the whole of StrictPreflight: a delegated root, this process
// placed in it or in its supervise child, and the root a real cgroup2
// directory. It returns the root opened without following a symlink.
func (m *Manager) strictRoot() (*os.File, error) {
	if m == nil || m.root == "" {
		return nil, ErrStrictUnsupported
	}
	own, e := os.ReadFile(selfCgroupFile)
	rel, ok := ParseSelfCgroup(string(own))
	if e != nil || !ok {
		return nil, ErrStrictUnsupported
	}
	current := filepath.Join("/sys/fs/cgroup", rel)
	rootPath := filepath.Clean(m.root)
	if current != rootPath && current != filepath.Join(rootPath, "supervise") {
		return nil, ErrStrictRefused
	}
	root, e := os.OpenFile(rootPath, os.O_RDONLY|unix.O_DIRECTORY|unix.O_NOFOLLOW, 0)
	if e != nil {
		return nil, ErrStrictUnsupported
	}
	var fs unix.Statfs_t
	if e = unix.Fstatfs(int(root.Fd()), &fs); e != nil || uint64(fs.Type) != uint64(unix.CGROUP2_SUPER_MAGIC) {
		root.Close()
		return nil, ErrStrictUnsupported
	}
	return root, nil
}

// StrictPreflight reports whether PrepareStrict's runtime checks that do
// not depend on the host pass now: the same function, not a copy, so a
// readiness check (moves S2) cannot approve a box PrepareStrict refuses
// for its placement or its mount. Nil means a strict spawn can be
// attempted; the per-host state (a populated group, a held reservation)
// is still PrepareStrict's to judge.
func (m *Manager) StrictPreflight() error {
	root, e := m.strictRoot()
	if e != nil {
		return e
	}
	return root.Close()
}

// PrepareStrict never uses accounting's fallback. Only the two existing tiers
// are accepted; a zero or missing limit cannot create an uncapped command.
func (m *Manager) PrepareStrict(host string, limit Limit) (*Prepared, error) {
	return m.prepareStrict(host, limit, false)
}

// prepareStrict is PrepareStrict and PrepareStrictParent: with genParent the
// group is pinned as a generation parent (memory.oom.group=0, controllers
// delegated to its children, Start refused) instead of a leaf.
func (m *Manager) prepareStrict(host string, limit Limit, genParent bool) (out *Prepared, err error) {
	if m == nil || m.root == "" {
		return nil, ErrStrictUnsupported
	}
	if !controlname.ValidHostShape(host) || (limit != TierDefault && limit != TierChildren) {
		return nil, ErrStrictRefused
	}
	root, e := m.strictRoot()
	if e != nil {
		return nil, e
	}
	p := &Prepared{strictPreparedState: &strictPreparedState{root: root, host: host, limits: map[string]*os.File{}, start: func(cmd *exec.Cmd) error { return cmd.Start() }}}
	reserved := false
	defer func() {
		if err != nil {
			p.closeFiles()
			if reserved {
				strictReservations.Lock()
				delete(strictReservations.held, p.identity)
				strictReservations.Unlock()
			}
		}
	}()
	// A group that took cgroup.kill is removed and made fresh first (amadan
	// #272). One a strict handle still holds stays, and the answer is the
	// reservation's own ErrStrictState.
	if e = m.freshen(host, m.hostDir(host)); e != nil {
		if errors.Is(e, ErrStrictState) {
			return nil, e
		}
		return nil, fmt.Errorf("%w: %w", ErrStrictRefused, e)
	}
	if e = unix.Mkdirat(int(root.Fd()), ExecName(host), 0755); e != nil && !errors.Is(e, unix.EEXIST) {
		return nil, ErrStrictRefused
	}
	p.dir, e = strictOpen(root, ExecName(host), unix.O_RDONLY|unix.O_DIRECTORY)
	if e != nil {
		return nil, ErrStrictRefused
	}
	p.identity, e = strictID(p.dir)
	if e != nil {
		return nil, ErrStrictRefused
	}
	strictReservations.Lock()
	if strictReservations.held[p.identity] {
		strictReservations.Unlock()
		return nil, ErrStrictState
	}
	strictReservations.held[p.identity] = true
	reserved = true
	strictReservations.Unlock()
	p.events, e = strictOpen(p.dir, "cgroup.events", unix.O_RDONLY)
	if e != nil {
		return nil, ErrStrictUnsupported
	}
	p.kill, e = strictOpen(p.dir, "cgroup.kill", unix.O_WRONLY)
	if e != nil {
		return nil, ErrStrictUnsupported
	}
	if populated, e := strictPopulation(p.events); e != nil || populated {
		return nil, ErrStrictRefused
	}
	oomGroup := "1"
	if genParent {
		oomGroup = "0"
	}
	p.expected = tierExpected(limit, oomGroup)
	for name := range p.expected {
		p.limits[name], e = strictOpen(p.dir, name, unix.O_RDWR)
		if e != nil {
			return nil, ErrStrictUnsupported
		}
	}
	for _, name := range []string{"memory.max", "memory.high", "memory.swap.max", "memory.oom.group"} {
		if e = strictWrite(p.limits[name], p.expected[name]); e != nil {
			return nil, ErrStrictRefused
		}
	}
	if e = p.limitsOK(); e != nil {
		return nil, e
	}
	if genParent {
		if e = delegateStrict(p.dir); e != nil {
			return nil, e
		}
		p.genParent, p.tier = true, limit
	}
	if e = p.identityOK(); e != nil {
		return nil, e
	}
	return p, nil
}

// Start consumes one attempt even if exec or containment fails. The caller
// supplies the complete command, including its approved executable descriptor.
func (p *Prepared) Start(cmd *exec.Cmd) error {
	if p == nil || p.strictPreparedState == nil {
		return ErrStrictState
	}
	if p.parent != nil {
		p.parent.mu.Lock()
		defer p.parent.mu.Unlock()
		if p.parent.closed || p.parent.empty || p.parent.draining.Load() {
			return ErrStrictState
		}
		if err := p.parent.identityOK(); err != nil {
			return err
		}
		if err := p.parent.limitsOK(); err != nil {
			return err
		}
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.genParent {
		// A generation parent holds no process (zero-downtime spec §7.3).
		return ErrStrictState
	}
	if p.closed || p.attempted || p.empty || p.draining.Load() {
		return ErrStrictState
	}
	p.attempted = true
	if cmd == nil || cmd.Process != nil || cmd.ProcessState != nil {
		return ErrStrictState
	}
	if e := p.identityOK(); e != nil {
		return e
	}
	if e := p.limitsOK(); e != nil {
		return e
	}
	if populated, e := strictPopulation(p.events); e != nil || populated {
		return ErrStrictRefused
	}
	setInto(cmd, int(p.dir.Fd()))
	cmd.SysProcAttr.Pdeathsig = syscall.SIGKILL
	if e := p.start(cmd); e != nil {
		if shouldLatch(e) {
			return ErrStrictUnsupported
		}
		return ErrStrictRefused
	}
	if cmd.Process == nil {
		return ErrStrictRefused
	}
	return nil
}

// Populated checks the retained group after the direct child has been reaped.
func (p *Prepared) Populated() (bool, error) {
	if p == nil || p.strictPreparedState == nil {
		return false, ErrStrictState
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if e := p.identityOK(); e != nil {
		return false, e
	}
	return strictPopulation(p.events)
}

// KillAndEmpty proves kernel population is zero before the caller's deadline.
// Failure retains the pinned cleanup handle and permits another bounded attempt.
func (p *Prepared) KillAndEmpty(ctx context.Context) error {
	if p == nil || p.strictPreparedState == nil {
		return ErrStrictState
	}
	p.draining.Store(true)
	if ctx == nil {
		return ErrStrictDeadline
	}
	deadline, ok := ctx.Deadline()
	if !ok || ctx.Err() != nil || !time.Now().Before(deadline) {
		return ErrStrictDeadline
	}
	lockTick := time.NewTicker(time.Millisecond)
	defer lockTick.Stop()
	for !p.mu.TryLock() {
		select {
		case <-ctx.Done():
			return ErrStrictDeadline
		case <-lockTick.C:
		}
	}
	defer p.mu.Unlock()
	if p.closed || p.empty {
		return ErrStrictState
	}
	if ctx.Err() != nil {
		return ErrStrictDeadline
	}
	if e := p.identityOK(); e != nil {
		return e
	}
	p.attempted = true
	// Marked first, so no later start reuses this directory (amadan #272,
	// recreate_linux.go).
	markKilled(p.identity)
	if e := strictWrite(p.kill, "1"); e != nil {
		return ErrStrictRefused
	}
	timer := time.NewTicker(5 * time.Millisecond)
	defer timer.Stop()
	for {
		if ctx.Err() != nil || !time.Now().Before(deadline) {
			return ErrStrictDeadline
		}
		if e := p.identityOK(); e != nil {
			return e
		}
		populated, e := strictPopulation(p.events)
		if e != nil {
			return ErrStrictRefused
		}
		if !populated {
			if ctx.Err() != nil || !time.Now().Before(deadline) {
				return ErrStrictDeadline
			}
			p.empty = true
			return nil
		}
		select {
		case <-ctx.Done():
			return ErrStrictDeadline
		case <-timer.C:
		}
	}
}

// Close never removes a path. A populated or unverifiable group retains its
// descriptors and reservation so the same handle can still attempt cleanup.
func (p *Prepared) Close() error {
	if p == nil || p.strictPreparedState == nil {
		return ErrStrictState
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	p.draining.Store(true)
	if p.closed {
		return ErrStrictState
	}
	if e := p.identityOK(); e != nil {
		return e
	}
	if populated, e := strictPopulation(p.events); e != nil || populated {
		return ErrStrictRefused
	}
	p.closed = true
	p.closeFiles()
	strictReservations.Lock()
	delete(strictReservations.held, p.identity)
	strictReservations.Unlock()
	return nil
}