//go:build linux

package cgroups

import (
	"errors"
	"fmt"
	"hash/crc32"
	"io"
	"math/rand"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	cgrps "github.com/containerd/cgroups/v3"
	"github.com/opencontainers/runtime-spec/specs-go"
	cgroupscfg "gitlab.com/gitlab-org/gitaly/v18/internal/gitaly/config/cgroups"
	"gitlab.com/gitlab-org/gitaly/v18/internal/log"
)

// cfs_period_us hardcoded to be 100ms.
const cfsPeriodUs uint64 = 100000

type cgroupHandler interface {
	setupParent(parentResources *specs.LinuxResources) error
	createCgroup(repoResources *specs.LinuxResources, cgroupPath string) error
	addToCgroup(pid int, cgroupPath string) error
	currentProcessCgroup() string
	repoPath(groupID int) string
	stats(path string) (CgroupStats, error)
	supportsCloneIntoCgroup() bool
}

type cgroupLock struct {
	created       atomic.Bool
	creationError error
	once          sync.Once
}

func (l *cgroupLock) isCreated() bool {
	return l.created.Load()
}

type cgroupStatus struct {
	m map[string]*cgroupLock
}

func newCgroupStatus(cfg cgroupscfg.Config, repoPath func(int) string) *cgroupStatus {
	status := &cgroupStatus{
		m: make(map[string]*cgroupLock, cfg.Repositories.Count),
	}

	// Pre-fill the map with all possible values.
	for i := 0; i < int(cfg.Repositories.Count); i++ {
		cgroupPath := repoPath(i)
		cgLock := &cgroupLock{}
		status.m[cgroupPath] = cgLock
	}

	return status
}

func (s *cgroupStatus) getLock(cgroupPath string) *cgroupLock {
	// We initialized all locks during construction.
	cgLock := s.m[cgroupPath]

	return cgLock
}

// randomizer is the interface of the Go random number generator.
type randomizer interface {
	// Intn returns a random integer in the range [0,n).
	Intn(n int) int
}

// CGroupManager is a manager class that implements specific methods related to cgroups
type CGroupManager struct {
	cfg           cgroupscfg.Config
	pid           int
	enabled       bool
	cgroupVersion uint8
	repoRes       *specs.LinuxResources
	status        *cgroupStatus
	handler       cgroupHandler
	rand          randomizer
}

func newCgroupManager(cfg cgroupscfg.Config, logger log.Logger, pid int) *CGroupManager {
	return newCgroupManagerWithMode(cfg, logger, pid, cgrps.Mode())
}

func newCgroupManagerWithMode(cfg cgroupscfg.Config, logger log.Logger, pid int, mode cgrps.CGMode) *CGroupManager {
	var handler cgroupHandler

	var cgroupVersion uint8
	switch mode {
	case cgrps.Legacy, cgrps.Hybrid:
		cgroupVersion = 1
		handler = newV1Handler(cfg, logger, pid)
	case cgrps.Unified:
		cgroupVersion = 2
		handler = newV2Handler(cfg, logger, pid)
		logger.Warn("Gitaly now includes experimental support for CgroupV2. Please proceed with caution and use this experimental feature at your own risk")
	default:
		logger.Warn("Gitaly has encountered an issue while trying to detect the version of the system's cgroup. As a result, all subsequent commands will be executed without cgroup support. Please check the system's cgroup configuration and try again")
		return nil
	}

	return &CGroupManager{
		cfg:           cfg,
		pid:           pid,
		cgroupVersion: cgroupVersion,
		handler:       handler,
		repoRes:       configRepositoryResources(cfg),
		status:        newCgroupStatus(cfg, handler.repoPath),
		rand:          rand.New(rand.NewSource(time.Now().UnixNano())),
	}
}

// Setup parent cgroups and repository sub cgroups
func (cgm *CGroupManager) Setup() error {
	if err := cgm.handler.setupParent(cgm.configParentResources()); err != nil {
		return err
	}

	cgroupPath := cgm.handler.repoPath(0)
	if err := cgm.maybeCreateCgroup(cgroupPath); err != nil {
		return fmt.Errorf("setup repository-scoped cgroup: %w", err)
	}

	cgm.enabled = true

	return nil
}

// Ready returns true if the Cgroup manager is configured and ready to use.
func (cgm *CGroupManager) Ready() bool {
	return cgm.enabled
}

// AddCommand adds a Cmd that has already started to a cgroup
func (cgm *CGroupManager) AddCommand(cmd *exec.Cmd, opts ...AddCommandOption) (string, error) {
	if cmd.Process == nil {
		return "", errors.New("cannot add command that has not yet been started")
	}

	cgroupPath := cgm.cgroupPathForCommand(cmd, opts)

	if err := cgm.maybeCreateCgroup(cgroupPath); err != nil {
		return "", fmt.Errorf("setup cgroup: %w", err)
	}

	return cgroupPath, cgm.handler.addToCgroup(cmd.Process.Pid, cgroupPath)
}

// SupportsCloneIntoCgroup returns whether this Manager supports the CloneIntoCgroup method.
// CloneIntoCgroup requires CLONE_INTO_CGROUP which is only supported with cgroup version 2
// with Linux 5.7 or newer.
func (cgm *CGroupManager) SupportsCloneIntoCgroup() bool {
	return cgm.handler.supportsCloneIntoCgroup()
}

// maybeCreateCgroup creates a cgroup if it yet hasn't been created.
func (cgm *CGroupManager) maybeCreateCgroup(cgroupPath string) error {
	lock := cgm.status.getLock(cgroupPath)

	lock.once.Do(func() {
		if err := cgm.handler.createCgroup(cgm.repoRes, cgroupPath); err != nil {
			lock.creationError = err
			return
		}

		lock.created.Store(true)
	})

	return lock.creationError
}

// CloneIntoCgroup configures the cgroup parameters UseCgroupFD and CgroupFD in SysProcAttr
// to start the command directly in the correct cgroup. On success, the function returns an io.Closer
// that must be closed after the command has been started to close the cgroup's file descriptor.
func (cgm *CGroupManager) CloneIntoCgroup(cmd *exec.Cmd, opts ...AddCommandOption) (string, io.Closer, error) {
	cgroupPath := cgm.cgroupPathForCommand(cmd, opts)

	if err := cgm.maybeCreateCgroup(cgroupPath); err != nil {
		return "", nil, fmt.Errorf("setup cgroup: %w", err)
	}

	cgroupDirPath := filepath.Join(cgm.cfg.Mountpoint, cgroupPath)
	dir, err := os.Open(cgroupDirPath)
	if err != nil {
		return "", nil, fmt.Errorf("open file: %w", err)
	}

	if cmd.SysProcAttr == nil {
		cmd.SysProcAttr = &syscall.SysProcAttr{}
	}

	cmd.SysProcAttr.UseCgroupFD = true
	cmd.SysProcAttr.CgroupFD = int(dir.Fd())

	return cgroupDirPath, dir, nil
}

// cgroupPathForCommand returns the path of the cgroup a given command should go in.
func (cgm *CGroupManager) cgroupPathForCommand(cmd *exec.Cmd, opts []AddCommandOption) string {
	var cfg addCommandCfg
	for _, opt := range opts {
		opt(&cfg)
	}

	key := cfg.cgroupKey
	if key == "" {
		key = strings.Join(cmd.Args, "/")
	}
	groupID := cgm.calcGroupID(cgm.rand, key, cgm.cfg.Repositories.Count, cgm.cfg.Repositories.MaxCgroupsPerRepo)
	return cgm.handler.repoPath(int(groupID))
}

func (cgm *CGroupManager) calcGroupID(rand randomizer, key string, count uint, allocationCount uint) uint {
	checksum := crc32.ChecksumIEEE([]byte(key))

	// Pick a starting point
	groupID := uint(checksum) % count
	if allocationCount <= 1 {
		return groupID
	}

	// Shift random distance [0, allocation) from the starting point. Wrap-around if needed.
	return (groupID + uint(rand.Intn(int(allocationCount)))) % count
}

// Stats returns cgroup accounting statistics collected by reading
// cgroupfs files.
func (cgm *CGroupManager) Stats() (Stats, error) {
	parentStats, err := cgm.handler.stats(cgm.currentProcessCgroup())
	if err != nil {
		return Stats{}, err
	}

	return Stats{
		ParentStats: parentStats,
	}, nil
}

func (cgm *CGroupManager) currentProcessCgroup() string {
	return cgm.handler.currentProcessCgroup()
}

func (cgm *CGroupManager) configParentResources() *specs.LinuxResources {
	cfsPeriodUs := cfsPeriodUs
	var parentResources specs.LinuxResources
	// Leave them `nil` so it takes kernel default unless cfg value above `0`.
	parentResources.CPU = &specs.LinuxCPU{}

	if cgm.cfg.CPUShares > 0 {
		parentResources.CPU.Shares = &cgm.cfg.CPUShares
	}

	if cgm.cfg.CPUQuotaUs > 0 {
		parentResources.CPU.Quota = &cgm.cfg.CPUQuotaUs
		parentResources.CPU.Period = &cfsPeriodUs
	}

	if cgm.cfg.MemoryBytes > 0 {
		parentResources.Memory = &specs.LinuxMemory{Limit: &cgm.cfg.MemoryBytes}
	}
	return &parentResources
}

func configRepositoryResources(cfg cgroupscfg.Config) *specs.LinuxResources {
	cfsPeriodUs := cfsPeriodUs
	var reposResources specs.LinuxResources
	// Leave them `nil` so it takes kernel default unless cfg value above `0`.
	reposResources.CPU = &specs.LinuxCPU{}

	if cfg.Repositories.CPUShares > 0 {
		reposResources.CPU.Shares = &cfg.Repositories.CPUShares
	}

	if cfg.Repositories.CPUQuotaUs > 0 {
		reposResources.CPU.Quota = &cfg.Repositories.CPUQuotaUs
		reposResources.CPU.Period = &cfsPeriodUs
	}

	if cfg.Repositories.MemoryBytes > 0 {
		reposResources.Memory = &specs.LinuxMemory{Limit: &cfg.Repositories.MemoryBytes}
	}
	return &reposResources
}

func pruneOldCgroups(cfg cgroupscfg.Config, logger log.Logger) {
	pruneOldCgroupsWithMode(cfg, logger, cgrps.Mode())
}

func pruneOldCgroupsWithMode(cfg cgroupscfg.Config, logger log.Logger, mode cgrps.CGMode) {
	if cfg.HierarchyRoot == "" {
		return
	}

	switch mode {
	case cgrps.Legacy, cgrps.Hybrid:
		pruneOldCgroupsV1(cfg, logger)
	case cgrps.Unified:
		pruneOldCgroupsV2(cfg, logger)
	default:
	}
}
