// Copyright (c) 2026 Storm Dragon
// SPDX-License-Identifier: GPL-3.0-or-later

package tentacle

import (
	"crypto/rand"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"syscall"
	"time"

	"git.stormux.org/storm/cthulhu/internal/input"
)

var ErrContainment = errors.New("tentacle resource containment failed")

// ResourceBudget includes the child and all its descendants. Tasks counts
// kernel threads as well as processes. CPUQuota is microseconds per 100 ms,
// so 100000 allows one CPU, independent of the machine's number of CPUs.
type ResourceBudget struct {
	CPUQuota    int64
	MemoryBytes int64
	Tasks       int64
}

type ResourceLimits struct {
	PerSession     ResourceBudget
	Total          ResourceBudget
	CleanupTimeout time.Duration
}

func DefaultResourceLimits() ResourceLimits {
	return ResourceLimits{
		PerSession:     ResourceBudget{CPUQuota: 25000, MemoryBytes: 256 << 20, Tasks: 64},
		Total:          ResourceBudget{CPUQuota: 100000, MemoryBytes: 1 << 30, Tasks: 256},
		CleanupTimeout: 2 * time.Second,
	}
}

func (limits ResourceLimits) validate() error {
	for _, budget := range []ResourceBudget{limits.PerSession, limits.Total} {
		if budget.CPUQuota < 1000 || budget.MemoryBytes < 1<<20 || budget.Tasks < 1 {
			return fmt.Errorf("%w: invalid resource budget", ErrContainment)
		}
	}
	if limits.Total.CPUQuota < limits.PerSession.CPUQuota ||
		limits.Total.MemoryBytes < limits.PerSession.MemoryBytes ||
		limits.Total.Tasks < limits.PerSession.Tasks || limits.CleanupTimeout <= 0 {
		return fmt.Errorf("%w: invalid aggregate budget or cleanup deadline", ErrContainment)
	}
	return nil
}

// NewContainedManager creates a private subtree in an already delegated cgroup
// v2 parent with cpu, memory and pids enabled in cgroup.subtree_control. It never
// moves the host or changes the caller's parent. Setup may perform filesystem
// I/O and belongs on a lifecycle worker, never an input callback.
//
// There is no fallback when delegation, controllers, cgroup.kill or atomic
// cgroup placement are unavailable. This is resource containment for trusted
// local code, not a security sandbox against same-user cgroup modification,
// migration, privileged scheduling, or external service launches.
func NewContainedManager(limits Limits, overrides []input.BindingOverride, parent string, resources ResourceLimits) (*Manager, error) {
	if err := resources.validate(); err != nil {
		return nil, err
	}
	manager, err := NewManager(limits, overrides)
	if err != nil {
		return nil, err
	}
	if !filepath.IsAbs(parent) {
		return nil, fmt.Errorf("%w: cgroup parent must be absolute", ErrContainment)
	}
	var stat syscall.Statfs_t
	if err := syscall.Statfs(parent, &stat); err != nil || stat.Type != 0x63677270 {
		return nil, fmt.Errorf("%w: parent is not a cgroup v2 filesystem", ErrContainment)
	}
	var identity [16]byte
	if _, err := rand.Read(identity[:]); err != nil {
		return nil, err
	}
	group, err := create_cgroup(filepath.Join(parent, fmt.Sprintf("cthulhu-tentacles-%x", identity)), resources.Total, limits.Processes)
	if err != nil {
		return nil, err
	}
	if err := write_control(group.path, "cgroup.subtree_control", "+cpu +memory +pids"); err != nil {
		return nil, errors.Join(err, group.remove())
	}
	manager.containment = &containment{group: group, limits: resources}
	return manager, nil
}

type containment struct {
	group  *cgroup
	limits ResourceLimits
}

type cgroup struct {
	path string
	file *os.File
}

func write_control(path, name, value string) error {
	// Do not create ordinary files if the kernel lacks a required controller.
	file, err := os.OpenFile(filepath.Join(path, name), os.O_WRONLY, 0)
	if err != nil {
		return fmt.Errorf("%w: open %s: %w", ErrContainment, name, err)
	}
	_, writeError := file.WriteString(value)
	return errors.Join(wrap_control_error(name, writeError), wrap_control_error(name, file.Close()))
}

func wrap_control_error(name string, err error) error {
	if err == nil {
		return nil
	}
	return fmt.Errorf("%w: %s: %w", ErrContainment, name, err)
}

func create_cgroup(path string, budget ResourceBudget, children int) (*cgroup, error) {
	if err := os.Mkdir(path, 0700); err != nil {
		return nil, fmt.Errorf("%w: create group: %w", ErrContainment, err)
	}
	group := &cgroup{path: path}
	settings := []struct{ name, value string }{
		{"cpu.max", strconv.FormatInt(budget.CPUQuota, 10) + " 100000"},
		{"memory.max", strconv.FormatInt(budget.MemoryBytes, 10)},
		{"memory.swap.max", "0"},
		{"pids.max", strconv.FormatInt(budget.Tasks, 10)},
		{"cgroup.max.descendants", strconv.Itoa(children)},
	}
	if children == 0 {
		// If any task exhausts memory, retire the entire session, not an
		// arbitrary worker while leaving a partially broken child alive.
		settings = append(settings, struct{ name, value string }{"memory.oom.group", "1"})
	}
	for _, setting := range settings {
		if err := write_control(path, setting.name, setting.value); err != nil {
			return nil, errors.Join(err, group.remove())
		}
	}
	// Check availability without issuing a kill before launch. Some kernels
	// SIGKILL CLONE_INTO_CGROUP children entering a previously killed group.
	killFile, err := os.OpenFile(filepath.Join(path, "cgroup.kill"), os.O_WRONLY, 0)
	if err != nil {
		return nil, errors.Join(wrap_control_error("open cgroup.kill", err), group.remove())
	}
	if err := killFile.Close(); err != nil {
		return nil, errors.Join(wrap_control_error("close cgroup.kill", err), group.remove())
	}
	file, err := os.Open(path)
	if err != nil {
		return nil, errors.Join(wrap_control_error("open group", err), group.remove())
	}
	group.file = file
	return group, nil
}

func (group *cgroup) remove() error {
	if group.file != nil {
		_ = group.file.Close()
		group.file = nil
	}
	return wrap_control_error("remove group", os.Remove(group.path))
}

// kill_and_remove waits only during teardown. Successful removal proves the
// kernel no longer has live tasks in the group, including setsid descendants.
func (group *cgroup) kill_and_remove(timeout time.Duration) error {
	defer func() {
		if group.file != nil {
			_ = group.file.Close()
			group.file = nil
		}
	}()
	if err := write_control(group.path, "cgroup.kill", "1"); err != nil {
		return err
	}
	deadline := time.NewTimer(timeout)
	defer deadline.Stop()
	ticker := time.NewTicker(10 * time.Millisecond)
	defer ticker.Stop()
	for {
		err := group.remove()
		if err == nil || !errors.Is(err, syscall.EBUSY) {
			return err
		}
		select {
		case <-ticker.C:
		case <-deadline.C:
			return fmt.Errorf("%w: cleanup deadline expired", ErrContainment)
		}
	}
}