// Copyright 2022 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 gctuner

import (
	"runtime"
	"runtime/debug"
	"testing"
	"time"

	"github.com/pingcap/failpoint"
	"github.com/pingcap/tidb/pkg/util/intest"
	"github.com/pingcap/tidb/pkg/util/memory"
	"github.com/stretchr/testify/require"
)

type mockAllocator struct {
	m [][]byte
}

func (a *mockAllocator) alloc(bytes int) (handle int) {
	sli := make([]byte, bytes)
	a.m = append(a.m, sli)
	return len(a.m) - 1
}

func (a *mockAllocator) free(handle int) {
	a.m[handle] = nil
}

func (a *mockAllocator) freeAll() {
	a.m = nil
	runtime.GC()
}

func TestGlobalMemoryTuner(t *testing.T) {
	require.True(t, intest.InTest)
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/util/gctuner/testMemoryLimitTuner", "return(true)"))
	defer func() {
		require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/util/gctuner/testMemoryLimitTuner"))
	}()
	// Close GOGCTuner
	gogcTuner := EnableGOGCTuner.Load()
	EnableGOGCTuner.Store(false)
	defer EnableGOGCTuner.Store(gogcTuner)

	memory.ServerMemoryLimit.Store(1 << 30)   // 1GB
	GlobalMemoryLimitTuner.SetPercentage(0.8) // 1GB * 80% = 800MB
	GlobalMemoryLimitTuner.UpdateMemoryLimit()
	require.True(t, GlobalMemoryLimitTuner.isValidValueSet.Load())
	defer func() {
		// If test.count > 1, wait tuning finished.
		require.Eventually(t, func() bool {
			//nolint: all_revive
			return GlobalMemoryLimitTuner.isValidValueSet.Load()
		}, 5*time.Second, 100*time.Millisecond)
		require.Eventually(t, func() bool {
			//nolint: all_revive
			return !GlobalMemoryLimitTuner.adjustPercentageInProgress.Load()
		}, 5*time.Second, 100*time.Millisecond)
		require.Eventually(t, func() bool {
			//nolint: all_revive
			return !GlobalMemoryLimitTuner.nextGCTriggeredByMemoryLimit.Load()
		}, 5*time.Second, 100*time.Millisecond)
	}()

	allocator := &mockAllocator{}
	defer allocator.freeAll()
	r := &runtime.MemStats{}
	getNowGCNum := func() uint32 {
		runtime.ReadMemStats(r)
		return r.NumGC
	}
	checkNextGCEqualMemoryLimit := func() {
		runtime.ReadMemStats(r)
		nextGC := r.NextGC
		memoryLimit := GlobalMemoryLimitTuner.calcMemoryLimit(GlobalMemoryLimitTuner.GetPercentage())
		// Refer to golang source code, nextGC = memoryLimit - nonHeapMemory - overageMemory - headroom
		require.True(t, nextGC < uint64(memoryLimit))
	}

	memory600mb := allocator.alloc(600 << 20)
	gcNum := getNowGCNum()

	memory210mb := allocator.alloc(210 << 20)
	require.Eventually(t, func() bool {
		return GlobalMemoryLimitTuner.adjustPercentageInProgress.Load() && gcNum < getNowGCNum()
	}, 5*time.Second, 100*time.Millisecond)
	// Test waiting for reset
	require.Eventually(t, func() bool {
		return GlobalMemoryLimitTuner.calcMemoryLimit(fallbackPercentage) == debug.SetMemoryLimit(-1)
	}, 5*time.Second, 100*time.Millisecond)
	gcNum = getNowGCNum()
	memory100mb := allocator.alloc(100 << 20)
	require.Eventually(t, func() bool {
		return gcNum == getNowGCNum()
	}, 5*time.Second, 100*time.Millisecond) // No GC

	allocator.free(memory210mb)
	allocator.free(memory100mb)
	runtime.GC()
	// Trigger GC in 80% again
	require.Eventually(t, func() bool {
		return GlobalMemoryLimitTuner.calcMemoryLimit(GlobalMemoryLimitTuner.GetPercentage()) == debug.SetMemoryLimit(-1)
	}, 5*time.Second, 100*time.Millisecond)
	time.Sleep(100 * time.Millisecond)
	gcNum = getNowGCNum()
	checkNextGCEqualMemoryLimit()
	memory210mb = allocator.alloc(210 << 20)
	require.Eventually(t, func() bool {
		return gcNum < getNowGCNum()
	}, 5*time.Second, 100*time.Millisecond)
	allocator.free(memory210mb)
	allocator.free(memory600mb)
}

func TestIssue48741(t *testing.T) {
	require.True(t, intest.InTest)
	// Close GOGCTuner
	gogcTuner := EnableGOGCTuner.Load()
	EnableGOGCTuner.Store(false)
	defer EnableGOGCTuner.Store(gogcTuner)

	// This test validates UpdateMemoryLimit behavior while the tuner is in "adjustPercentageInProgress" mode.
	// Historically, the previous implementation tried to reproduce it by allocating huge memory and waiting for GC
	// callbacks (finalizer) to kick in, which is flaky in CI due to scheduling variance. Here we simulate the
	// in-progress state directly, which is deterministic and still validates the contract.

	// Snapshot global state to avoid impacting other tests.
	prevServerMemoryLimit := memory.ServerMemoryLimit.Load()
	prevPercentage := GlobalMemoryLimitTuner.GetPercentage()
	prevAdjustInProgress := GlobalMemoryLimitTuner.adjustPercentageInProgress.Load()
	prevBeforeAdjustMem := GlobalMemoryLimitTuner.serverMemLimitBeforeAdjust.Load()
	prevBeforeAdjustPerc := GlobalMemoryLimitTuner.percentageBeforeAdjust.Load()
	prevValidValueSet := GlobalMemoryLimitTuner.isValidValueSet.Load()
	prevAdjustDisabled := GlobalMemoryLimitTuner.adjustDisabled.Load()
	prevGoMemLimit := debug.SetMemoryLimit(-1)
	defer func() {
		memory.ServerMemoryLimit.Store(prevServerMemoryLimit)
		GlobalMemoryLimitTuner.SetPercentage(prevPercentage)
		GlobalMemoryLimitTuner.serverMemLimitBeforeAdjust.Store(prevBeforeAdjustMem)
		GlobalMemoryLimitTuner.percentageBeforeAdjust.Store(prevBeforeAdjustPerc)
		GlobalMemoryLimitTuner.isValidValueSet.Store(prevValidValueSet)
		GlobalMemoryLimitTuner.adjustDisabled.Store(prevAdjustDisabled)
		GlobalMemoryLimitTuner.adjustPercentageInProgress.Store(prevAdjustInProgress)
		debug.SetMemoryLimit(prevGoMemLimit)
	}()

	// Use smaller memory footprints to avoid stressing CI runners.
	const baseMem uint64 = 256 << 20 // 256MiB
	const newMem uint64 = 384 << 20  // 384MiB

	// Ensure UpdateMemoryLimit will take effect.
	GlobalMemoryLimitTuner.adjustDisabled.Store(0)
	GlobalMemoryLimitTuner.isValidValueSet.Store(true)

	memory.ServerMemoryLimit.Store(baseMem)
	GlobalMemoryLimitTuner.SetPercentage(0.8)
	GlobalMemoryLimitTuner.UpdateMemoryLimit()

	// Simulate "tunning in progress": memory limit temporarily set to fallbackPercentage.
	GlobalMemoryLimitTuner.adjustPercentageInProgress.Store(true)
	GlobalMemoryLimitTuner.serverMemLimitBeforeAdjust.Store(baseMem)
	GlobalMemoryLimitTuner.percentageBeforeAdjust.Store(GlobalMemoryLimitTuner.GetPercentage())
	debug.SetMemoryLimit(GlobalMemoryLimitTuner.calcMemoryLimit(fallbackPercentage))

	// 1) Not modified: calling UpdateMemoryLimit should not override the fallback when inputs are unchanged.
	GlobalMemoryLimitTuner.UpdateMemoryLimit()
	require.Equal(t, GlobalMemoryLimitTuner.calcMemoryLimit(fallbackPercentage), debug.SetMemoryLimit(-1))

	// 2) Modified: changing ServerMemoryLimit should be reflected immediately even during tunning.
	memory.ServerMemoryLimit.Store(newMem)
	GlobalMemoryLimitTuner.UpdateMemoryLimit()
	require.Equal(t, GlobalMemoryLimitTuner.calcMemoryLimit(GlobalMemoryLimitTuner.GetPercentage()), debug.SetMemoryLimit(-1))
}

func TestSetMemoryLimit(t *testing.T) {
	GlobalMemoryLimitTuner.DisableAdjustMemoryLimit()
	memory.ServerMemoryLimit.Store(1 << 30)   // 1GB
	GlobalMemoryLimitTuner.SetPercentage(0.8) // 1GB * 80% = 800MB
	GlobalMemoryLimitTuner.UpdateMemoryLimit()
	require.Equal(t, initGOMemoryLimitValue, debug.SetMemoryLimit(-1))
	GlobalMemoryLimitTuner.EnableAdjustMemoryLimit()
	GlobalMemoryLimitTuner.UpdateMemoryLimit()
	require.Equal(t, int64(1<<30*80/100), debug.SetMemoryLimit(-1))
}