// Copyright 2023 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package serverless

import (
	"bytes"
	"encoding/json"
	"net/http"
	"os"
	"runtime"
	"strconv"
	"time"

	"github.com/pingcap/tidb/pkg/metrics"
	"github.com/pingcap/tidb/pkg/util/gctuner"
	"github.com/pingcap/tidb/pkg/util/logutil"
	"github.com/pingcap/tidb/pkg/util/memory"
	"go.uber.org/zap"
)

var (
	namespace      string
	podName        string
	nodeIP         string
	minMemoryLimit uint64
	maxMemoryLimit uint64

	// Serverless memory scaler knobs (env-configurable).
	//
	// Env:
	// - DISABLE_MEMORY_SHRINK=true: disable scale-down to reduce low-load GC jitter.
	// - MEMORY_SHRINK_COOLDOWN=5m: minimum time since last scale-up before scale-down is allowed.
	//   Non-positive duration (e.g. 0s) disables shrink as an operator-friendly shortcut.
	disableMemoryShrink  bool
	memoryShrinkCooldown = time.Minute
)

// StartMemoryScaler starts a memory scaler.
func StartMemoryScaler(q chan struct{}) {
	namespace = os.Getenv("NAMESPACE")
	podName = os.Getenv("POD_NAME")
	nodeIP = os.Getenv("NODE_IP")
	minMemoryLimitStr := os.Getenv("MIN_MEMORY_LIMIT")
	maxMemoryLimitStr := os.Getenv("MAX_MEMORY_LIMIT")
	disableMemoryShrinkStr := os.Getenv("DISABLE_MEMORY_SHRINK")
	memoryShrinkCooldownStr := os.Getenv("MEMORY_SHRINK_COOLDOWN")

	if minMemoryLimitStr == "" || maxMemoryLimitStr == "" {
		logutil.BgLogger().Info("env MIN_MEMORY_LIMIT or MAX_MEMORY_LIMIT not set, skip init memory limits")
		return
	}

	if disableMemoryShrinkStr != "" {
		if v, err := strconv.ParseBool(disableMemoryShrinkStr); err != nil {
			logutil.BgLogger().Warn("env DISABLE_MEMORY_SHRINK is bad format, ignore it",
				zap.String("value", disableMemoryShrinkStr),
				zap.Error(err),
			)
		} else {
			disableMemoryShrink = v
		}
	}
	if memoryShrinkCooldownStr != "" {
		if d, err := time.ParseDuration(memoryShrinkCooldownStr); err != nil {
			logutil.BgLogger().Warn("env MEMORY_SHRINK_COOLDOWN is bad format, ignore it",
				zap.String("value", memoryShrinkCooldownStr),
				zap.Error(err),
			)
		} else if d <= 0 {
			// Treat non-positive value as disabling shrink (operator-friendly).
			disableMemoryShrink = true
		} else {
			memoryShrinkCooldown = d
		}
	}

	var err error
	minMemoryLimit, err = strconv.ParseUint(minMemoryLimitStr, 10, 64)
	if err != nil {
		logutil.BgLogger().Info("env MIN_MEMORY_LIMIT is bad format, skip init memory limits")
		return
	}

	maxMemoryLimit, err = strconv.ParseUint(maxMemoryLimitStr, 10, 64)
	if err != nil {
		logutil.BgLogger().Info("env MAX_MEMORY_LIMIT is bad format, skip init memory limits")
		return
	}

	memory.ServerMemoryLimitOriginText.Store(minMemoryLimitStr)
	memory.ServerMemoryLimit.Store(minMemoryLimit)
	memory.MaxServerMemoryLimit.Store(maxMemoryLimit)
	memory.MaxServerMemoryLimitText.Store(maxMemoryLimitStr)
	gctuner.GlobalMemoryLimitTuner.UpdateMemoryLimit()

	if namespace == "" || podName == "" || nodeIP == "" {
		logutil.BgLogger().Info("env NAMESPACE, POD_NAME, NODE_IP not set, skip start memory scaler loop",
			zap.Bool("disable_memory_shrink", disableMemoryShrink),
			zap.Duration("memory_shrink_cooldown", memoryShrinkCooldown),
		)
		return
	}

	go NewMemoryScaler(q).Run()
	logutil.BgLogger().Info("start memory scaler success",
		zap.Bool("disable_memory_shrink", disableMemoryShrink),
		zap.Duration("memory_shrink_cooldown", memoryShrinkCooldown),
	)
}

// MemoryScaler is used to scale memory limit of a pod.
type MemoryScaler struct {
	exitCh chan struct{}
}

// NewMemoryScaler creates a new MemoryScaleHandle.
func NewMemoryScaler(exitCh chan struct{}) *MemoryScaler {
	return &MemoryScaler{
		exitCh: exitCh,
	}
}

// Pod stores memory usage information of a pod.
type Pod struct {
	Name      string `json:"name,omitempty"`
	Namespace string `json:"namespace,omitempty"`
	Limit     int64  `json:"limit,omitempty"`
	MinLimit  int64  `json:"min_limit,omitempty"`
	MaxLimit  int64  `json:"max_limit,omitempty"`
	Used      int64  `json:"used,omitempty"`
}

func (ms *MemoryScaler) report(stat *runtime.MemStats, limit uint64) {
	pod := &Pod{
		Name:      podName,
		Namespace: namespace,
		Limit:     int64(limit),
		MinLimit:  int64(minMemoryLimit),
		MaxLimit:  int64(maxMemoryLimit),
		Used:      int64(stat.HeapInuse),
	}
	data, _ := json.Marshal(pod)
	res, err := http.Post("http://"+nodeIP+":4040/report", "application/json", bytes.NewReader(data))
	if err != nil {
		logutil.BgLogger().Error("report memory usage failed", zap.Error(err))
		metrics.VPAReportCounter.WithLabelValues("failed").Inc()
		return
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		logutil.BgLogger().Error("report memory usage failed", zap.Int("status", res.StatusCode))
		metrics.VPAReportCounter.WithLabelValues("failed").Inc()
		return
	}
	metrics.VPAReportCounter.WithLabelValues("ok").Inc()
}

// ScaleRequest is the request to scale a pod's memory limit.
type ScaleRequest struct {
	Namespace string `json:"namespace,omitempty"`
	Name      string `json:"name,omitempty"`
	Limit     int64  `json:"limit,omitempty"`
}

func (ms *MemoryScaler) scaleMemory(to uint64) bool {
	logutil.BgLogger().Info("begin to scale memory", zap.Uint64("new_limit", to))
	req := &ScaleRequest{
		Namespace: namespace,
		Name:      podName,
		Limit:     int64(to),
	}
	data, _ := json.Marshal(req)
	client := &http.Client{Timeout: time.Second}
	res, err := client.Post("http://"+nodeIP+":4040/scale", "application/json", bytes.NewReader(data))
	if err != nil {
		logutil.BgLogger().Error("scale memory failed", zap.Error(err))
		metrics.VPAScaleMemoryCounter.WithLabelValues("failed").Inc()
		return false
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		logutil.BgLogger().Error("scale memory failed", zap.Int("status", res.StatusCode))
		metrics.VPAScaleMemoryCounter.WithLabelValues("failed").Inc()
		return false
	}
	logutil.BgLogger().Info("scale memory success", zap.Uint64("limit", to))
	metrics.VPAScaleMemoryCounter.WithLabelValues("ok").Inc()
	memory.ServerMemoryLimit.Store(to)
	gctuner.GlobalMemoryLimitTuner.UpdateMemoryLimit()
	return true
}

const maxMemoryScale = 2 * 1024 * 1024 * 1024 // 2GB

// Run runs the memory scale handle if it is configured.
func (ms *MemoryScaler) Run() {
	if minMemoryLimit == 0 || maxMemoryLimit == 0 {
		return
	}

	ticker := time.NewTicker(time.Millisecond * 100)
	defer ticker.Stop()
	lastReport, lastIncrease, lastIncreseFailed := time.Now(), time.Now(), time.Now()
	for {
		select {
		case <-ticker.C:
			stats := memory.ForceReadMemStats()
			limit := memory.ServerMemoryLimit.Load()
			metrics.VPAMemoryGauge.WithLabelValues("heap_inuse").Set(float64(stats.HeapInuse))
			metrics.VPAMemoryGauge.WithLabelValues("sys").Set(float64(stats.Sys))
			metrics.VPAMemoryGauge.WithLabelValues("heap_released").Set(float64(stats.HeapReleased))
			metrics.VPAMemoryGauge.WithLabelValues("limit").Set(float64(limit))

			if time.Since(lastReport) > time.Second {
				ms.report(stats, limit)
				lastReport = time.Now()
			}
			if limit != maxMemoryLimit && stats.HeapInuse > limit*8/10 && time.Since(lastIncreseFailed) > time.Second {
				if limit >= maxMemoryScale {
					limit += maxMemoryScale
				} else {
					limit *= 2
				}
				if limit > maxMemoryLimit {
					limit = maxMemoryLimit
				}
				ok := ms.scaleMemory(limit)
				if ok {
					lastIncrease = time.Now()
				} else {
					lastIncreseFailed = time.Now()
				}
			} else if !disableMemoryShrink &&
				limit != minMemoryLimit &&
				time.Since(lastIncrease) > memoryShrinkCooldown &&
				stats.Sys-stats.HeapReleased < limit/2 {
				limit = limit * 2 / 3
				if limit < minMemoryLimit {
					limit = minMemoryLimit
				}
				ms.scaleMemory(limit)
			}
		case <-ms.exitCh:
			metrics.VPAMemoryGauge.WithLabelValues("heap_inuse").Set(float64(0))
			metrics.VPAMemoryGauge.WithLabelValues("sys").Set(0)
			metrics.VPAMemoryGauge.WithLabelValues("heap_released").Set(0)
			metrics.VPAMemoryGauge.WithLabelValues("limit").Set(0)
			return
		}
	}
}