// Copyright 2024 TiKV Project Authors.
//
// 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 rule

import (
	"context"
	"encoding/hex"
	"encoding/json"
	"os"
	"path/filepath"
	"strconv"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"github.com/tikv/pd/pkg/core"
	"github.com/tikv/pd/pkg/keyspace"
	keyspaceconstant "github.com/tikv/pd/pkg/keyspace/constant"
	"github.com/tikv/pd/pkg/mock/mockcluster"
	"github.com/tikv/pd/pkg/mock/mockconfig"
	"github.com/tikv/pd/pkg/schedule/checker"
	"github.com/tikv/pd/pkg/schedule/labeler"
	"github.com/tikv/pd/pkg/storage/endpoint"
	"github.com/tikv/pd/pkg/storage/kv"
	"github.com/tikv/pd/pkg/utils/etcdutil"
	"github.com/tikv/pd/pkg/utils/keypath"
	clientv3 "go.etcd.io/etcd/client/v3"
	"go.etcd.io/etcd/server/v3/embed"
)

const (
	clusterID = uint64(20240117)
	rulesNum  = 16384
)

func TestLoadLargeRules(t *testing.T) {
	re := require.New(t)
	ctx, client, clean := prepare(t)
	defer clean()
	runWatcherLoadLabelRule(ctx, re, client)
}

func TestWatchCanonicalKeyspaceRule(t *testing.T) {
	re := require.New(t)
	ctx, client, clean := prepareEtcd(t)
	defer clean()

	const id = uint32(42)
	idText := strconv.FormatUint(uint64(id), endpoint.SpaceIDBase)
	ruleKey := keypath.RegionLabelPathPrefix() + "/" + keyspaceconstant.RegionLabelIDPrefix + idText
	initialRule := makeCanonicalKeyspaceRule(id, false)
	value, err := json.Marshal(initialRule)
	re.NoError(err)
	_, err = clientv3.NewKV(client).Put(ctx, ruleKey, string(value))
	re.NoError(err)

	storage := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
	labelerManager, err := labeler.NewRegionLabeler(ctx, storage, time.Hour)
	re.NoError(err)
	defer labelerManager.Close()

	checkerCtx, cancelChecker := context.WithCancel(ctx)
	opts := mockconfig.NewTestOptions()
	cluster := mockcluster.NewCluster(checkerCtx, opts)
	defer func() {
		cancelChecker()
		cluster.RegionLabeler.Close()
	}()
	checkerController := checker.NewController(
		checkerCtx,
		cluster,
		opts,
		cluster.GetRuleManager(),
		labelerManager,
		nil,
	)

	watchCtx, cancelWatch := context.WithCancel(ctx)
	rw := &Watcher{
		ctx:                   watchCtx,
		cancel:                cancelWatch,
		rulesPathPrefix:       keypath.RulesPathPrefix(),
		ruleCommonPathPrefix:  keypath.RuleCommonPathPrefix(),
		ruleGroupPathPrefix:   keypath.RuleGroupPathPrefix(),
		regionLabelPathPrefix: keypath.RegionLabelPathPrefix(),
		etcdClient:            client,
		ruleStorage:           storage,
		checkerController:     checkerController,
		regionLabeler:         labelerManager,
	}
	defer rw.Close()
	re.NoError(rw.initializeRegionLabelWatcher())

	bound := keyspace.MakeRegionBound(id)
	rawRegion := core.NewTestRegionInfo(1, 1, bound.RawLeftBound, bound.RawRightBound)
	txnRegion := core.NewTestRegionInfo(2, 1, bound.TxnLeftBound, bound.TxnRightBound)
	re.Equal(idText, labelerManager.GetRegionLabel(rawRegion, keyspaceconstant.RegionLabelKey))
	re.Equal(idText, labelerManager.GetRegionLabel(txnRegion, keyspaceconstant.RegionLabelKey))
	re.Equal([][]byte{
		bound.RawLeftBound,
		bound.RawRightBound,
		bound.TxnLeftBound,
		bound.TxnRightBound,
	}, labelerManager.GetSplitKeys(nil, nil))

	txnOnlyRule := makeCanonicalKeyspaceRule(id, true)
	value, err = json.Marshal(txnOnlyRule)
	re.NoError(err)
	_, err = clientv3.NewKV(client).Put(ctx, ruleKey, string(value))
	re.NoError(err)
	re.Eventually(func() bool {
		return labelerManager.GetRegionLabel(rawRegion, keyspaceconstant.RegionLabelKey) == "" &&
			labelerManager.GetRegionLabel(txnRegion, keyspaceconstant.RegionLabelKey) == idText
	}, 10*time.Second, 10*time.Millisecond)
	re.Equal([][]byte{bound.TxnLeftBound, bound.TxnRightBound}, labelerManager.GetSplitKeys(nil, nil))

	_, err = clientv3.NewKV(client).Delete(ctx, ruleKey)
	re.NoError(err)
	re.Eventually(func() bool {
		return labelerManager.GetRegionLabel(txnRegion, keyspaceconstant.RegionLabelKey) == "" &&
			len(labelerManager.GetSplitKeys(nil, nil)) == 0
	}, 10*time.Second, 10*time.Millisecond)
}

func BenchmarkLoadLargeRules(b *testing.B) {
	re := require.New(b)
	ctx, client, clean := prepare(b)
	defer clean()

	b.ResetTimer() // Resets the timer to ignore initialization time in the benchmark

	for n := 0; n < b.N; n++ {
		runWatcherLoadLabelRule(ctx, re, client)
	}
}

func runWatcherLoadLabelRule(ctx context.Context, re *require.Assertions, client *clientv3.Client) {
	storage := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
	labelerManager, err := labeler.NewRegionLabeler(ctx, storage, time.Hour)
	re.NoError(err)
	defer labelerManager.Close()
	ctx, cancel := context.WithCancel(ctx)
	rw := &Watcher{
		ctx:                   ctx,
		cancel:                cancel,
		rulesPathPrefix:       keypath.RulesPathPrefix(),
		ruleCommonPathPrefix:  keypath.RuleCommonPathPrefix(),
		ruleGroupPathPrefix:   keypath.RuleGroupPathPrefix(),
		regionLabelPathPrefix: keypath.RegionLabelPathPrefix(),
		etcdClient:            client,
		ruleStorage:           storage,
		regionLabeler:         labelerManager,
	}
	defer rw.Close()
	err = rw.initializeRegionLabelWatcher()
	re.NoError(err)
	allRules, err := labelerManager.GetAllLabelRules()
	re.NoError(err)
	re.Len(allRules, rulesNum)
}

func prepare(t require.TestingT) (context.Context, *clientv3.Client, func()) {
	re := require.New(t)
	ctx, client, clean := prepareEtcd(t)

	for i := 1; i < rulesNum+1; i++ {
		regionBound := keyspace.MakeRegionBound(uint32(i))
		rule := &labeler.LabelRule{
			ID:       "test_" + strconv.Itoa(i),
			Labels:   []labeler.RegionLabel{{Key: "test", Value: "test"}},
			RuleType: labeler.KeyRange,
			Data: labeler.MakeKeyRanges(
				hex.EncodeToString(regionBound.TxnLeftBound),
				hex.EncodeToString(regionBound.TxnRightBound),
				hex.EncodeToString(regionBound.RawLeftBound),
				hex.EncodeToString(regionBound.RawRightBound),
			),
		}
		value, err := json.Marshal(rule)
		re.NoError(err)
		key := keypath.RegionLabelPathPrefix() + "/" + rule.ID
		_, err = clientv3.NewKV(client).Put(ctx, key, string(value))
		re.NoError(err)
	}

	return ctx, client, clean
}

func prepareEtcd(t require.TestingT) (context.Context, *clientv3.Client, func()) {
	re := require.New(t)
	ctx, cancel := context.WithCancel(context.Background())
	cfg := etcdutil.NewTestSingleConfig()
	cfg.Dir = filepath.Join(os.TempDir(), "/pd_tests")
	os.RemoveAll(cfg.Dir)
	etcd, err := embed.StartEtcd(cfg)
	re.NoError(err)
	client, err := etcdutil.CreateEtcdClient(nil, cfg.ListenClientUrls)
	re.NoError(err)
	<-etcd.Server.ReadyNotify()

	return ctx, client, func() {
		cancel()
		client.Close()
		etcd.Close()
		os.RemoveAll(cfg.Dir)
	}
}

func makeCanonicalKeyspaceRule(id uint32, skipRaw bool) *labeler.LabelRule {
	bound := keyspace.MakeRegionBound(id)
	data := labeler.MakeKeyRanges(
		hex.EncodeToString(bound.TxnLeftBound),
		hex.EncodeToString(bound.TxnRightBound),
	)
	if !skipRaw {
		data = labeler.MakeKeyRanges(
			hex.EncodeToString(bound.RawLeftBound),
			hex.EncodeToString(bound.RawRightBound),
			hex.EncodeToString(bound.TxnLeftBound),
			hex.EncodeToString(bound.TxnRightBound),
		)
	}
	idText := strconv.FormatUint(uint64(id), endpoint.SpaceIDBase)
	return &labeler.LabelRule{
		ID:       keyspaceconstant.RegionLabelIDPrefix + idText,
		Index:    0,
		Labels:   []labeler.RegionLabel{{Key: keyspaceconstant.RegionLabelKey, Value: idText}},
		RuleType: labeler.KeyRange,
		Data:     data,
	}
}