package memlimit

import (
	"os"
	"runtime/debug"
	"strconv"
	"strings"

	"github.com/pingcap/dbaas/billing2/utils/log"
	"go.uber.org/zap"
)

const (
	// cgroupMemoryLimitRatio is the fraction of the cgroup memory limit
	// allocated to Go's soft memory limit. The remaining ~10% is reserved
	// for goroutine stacks, runtime metadata, OS thread stacks, and other
	// non-heap overhead.
	cgroupMemoryLimitRatio = 0.9

	// maxPlausibleLimit filters out "unlimited" values during cgroup file
	// reading. cgroup v1 reports ~math.MaxInt64 when no limit is set.
	maxPlausibleLimit int64 = 1 << 40 // 1 TiB

	// minGoMemLimit is the floor for the calculated soft memory limit.
	// Only applied when the cgroup limit itself is above this value, so
	// the floor never exceeds the real container ceiling.
	minGoMemLimit int64 = 900 << 20 // 900 MiB
)

var (
	// cgroup file paths, overridable in tests
	detectV2Path = "/sys/fs/cgroup/memory.max"
	detectV1Path = "/sys/fs/cgroup/memory/memory.limit_in_bytes"
)

// SetupFromCgroup detects the container's cgroup memory limit and sets
// Go's soft memory limit (GOMEMLIMIT) to 90% of it.
//
// This makes the garbage collector aware of the container memory ceiling,
// preventing unbounded heap growth from triggering cgroup OOM kills.
//
// Compatibility:
//   - Respects an explicit GOMEMLIMIT env var (does not override).
//   - Supports both cgroup v1 and v2.
//   - No-op outside containers (bare metal, local dev) — falls back to
//     Go's default (no soft limit).
//   - Automatically adapts when the container memory limit changes
//     (e.g., machine upgrade from 1 Gi → 2 Gi) — no code or config change
//     needed.
//
// Must be called once at process startup, after logger init and before
// any significant memory allocation.
func SetupFromCgroup() {
	if os.Getenv("GOMEMLIMIT") != "" {
		log.Info("GOMEMLIMIT already set via environment variable, skipping cgroup auto-detection")
		return
	}

	limit, source, ok := detect()
	if !ok {
		log.Info("cgroup memory limit not detected, using Go default GOMEMLIMIT")
		return
	}

	goMemLimit := int64(float64(limit) * cgroupMemoryLimitRatio)

	if goMemLimit < minGoMemLimit && limit > minGoMemLimit {
		log.Info("cgroup limit is close to floor, clamping GOMEMLIMIT",
			zap.String("source", source),
			zap.Int64("cgroupLimitBytes", limit),
			zap.Int64("calculatedGoMemLimit", goMemLimit),
			zap.Int64("floor", minGoMemLimit))
		goMemLimit = minGoMemLimit
	}

	previous := debug.SetMemoryLimit(goMemLimit)
	log.Info("set Go soft memory limit from cgroup",
		zap.String("source", source),
		zap.Int64("cgroupLimitBytes", limit),
		zap.Int64("goMemLimitBytes", goMemLimit),
		zap.Int64("previousLimitBytes", previous))
}

func detect() (limit int64, source string, ok bool) {
	if v, found := readLimit(detectV2Path); found {
		return v, "cgroup.v2", true
	}
	if v, found := readLimit(detectV1Path); found {
		return v, "cgroup.v1", true
	}
	return 0, "", false
}

func readLimit(path string) (int64, bool) {
	data, err := os.ReadFile(path)
	if err != nil {
		return 0, false
	}

	value := strings.TrimSpace(string(data))
	if value == "" || value == "max" {
		return 0, false
	}

	limit, err := strconv.ParseInt(value, 10, 64)
	if err != nil || limit <= 0 {
		return 0, false
	}

	if limit > maxPlausibleLimit {
		return 0, false
	}

	return limit, true
}