// Copyright 2021 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 labeler

import (
	"context"
	"errors"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pingcap/failpoint"
	"github.com/pingcap/log"
	"github.com/tikv/pd/pkg/core"
	"github.com/tikv/pd/pkg/errs"
	"github.com/tikv/pd/pkg/storage/endpoint"
	"github.com/tikv/pd/pkg/storage/kv"
	"github.com/tikv/pd/pkg/utils/logutil"
	"github.com/tikv/pd/pkg/utils/syncutil"
	"go.uber.org/zap"
)

// RegionLabeler is utility to label regions.
type RegionLabeler struct {
	storage   endpoint.RuleStorage
	mu        syncutil.RWMutex
	ruleIndex labelRuleIndex
	ctx       context.Context
	isReady   atomic.Bool
	wg        *sync.WaitGroup
	cancel    context.CancelFunc
}

// Lock acquires the labeler's write lock.
func (l *RegionLabeler) Lock() {
	l.mu.Lock()
}

// Unlock publishes all derived rule-index updates before releasing the write
// lock. This keeps a dirty index private to the write-side critical section.
func (l *RegionLabeler) Unlock() {
	defer l.mu.Unlock()
	l.ruleIndex.buildRanges()
}

// RLock acquires the labeler's read lock.
func (l *RegionLabeler) RLock() {
	l.mu.RLock()
}

// RUnlock releases the labeler's read lock.
func (l *RegionLabeler) RUnlock() {
	l.mu.RUnlock()
}

// NewRegionLabeler creates a Labeler instance.
func NewRegionLabeler(parentCtx context.Context, storage endpoint.RuleStorage, gcInterval time.Duration, asyncLoad ...bool) (*RegionLabeler, error) {
	start := time.Now()
	defer func() {
		newRegionLabelerDuration.Observe(time.Since(start).Seconds())
	}()

	ctx, cancel := context.WithCancel(parentCtx)

	l := &RegionLabeler{
		storage:   storage,
		ruleIndex: newLabelRuleIndex(nil),
		ctx:       ctx,
		wg:        &sync.WaitGroup{},
		cancel:    cancel,
	}

	if len(asyncLoad) > 0 && asyncLoad[0] {
		l.wg.Add(1)
		go func() {
			defer l.wg.Done()
			failpoint.Inject("pauseLoadRules", func(val failpoint.Value) {
				if s, ok := val.(string); ok {
					log.Info("skip loading region labeler rules", zap.String("duration", s))
					expected, err := time.ParseDuration(s)
					if err != nil {
						log.Error("failed to parse duration", zap.String("duration", s), zap.Error(err))
						return
					}
					select {
					case <-time.After(expected):
					case <-ctx.Done():
						return
					}
				}
			})
			// if success is false, it means the context is canceled, and we should exit the goroutine.
			if success := l.mustLoadRules(); !success {
				log.Warn("failed to load region labeler rules, exit")
				return
			}
			l.wg.Add(1)
			go l.doGC(gcInterval)
		}()
	} else {
		if err := l.loadRules(); err != nil {
			cancel()
			return nil, err
		}
		l.isReady.Store(true)
		log.Info("new region labeler created", zap.Int("label-rules-count", len(l.ruleIndex.rules)))
		l.wg.Add(1)
		go l.doGC(gcInterval)
	}
	return l, nil
}

// IsReady returns whether the labeler has loaded all rules.
//
// A canceled labeler is never ready, even if isReady was already set:
// l.ctx's parent can be canceled directly by callers such as
// RaftCluster.Stop (which cancels its own context before calling
// labeler.Close), bypassing any lock this labeler holds. So readiness is
// re-validated against the context on every call instead of relying solely
// on a flag that a concurrent, uncontrolled cancellation could race past.
func (l *RegionLabeler) IsReady() bool {
	if !l.isReady.Load() {
		return false
	}
	select {
	case <-l.ctx.Done():
		return false
	default:
		return true
	}
}

// doGC's caller must call l.wg.Add(1) first; Close()'s wg.Wait() only
// covers goroutines tracked this way, and doGC accesses labeler storage
// that Close()'s callers may tear down right after Close() returns.
func (l *RegionLabeler) doGC(gcInterval time.Duration) {
	defer l.wg.Done()
	defer logutil.LogPanic()

	ticker := time.NewTicker(gcInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			l.checkAndClearExpiredLabels()
			log.Debug("region labeler GC")
		case <-l.ctx.Done():
			log.Info("region labeler GC stopped")
			return
		}
	}
}

func (l *RegionLabeler) checkAndClearExpiredLabels() {
	now := time.Now()
	l.Lock()
	defer l.Unlock()

	if l.ruleIndex.minExpire == nil || l.ruleIndex.minExpire.After(now) {
		return
	}
	var err error

	for key, rule := range l.ruleIndex.rules {
		if !l.ruleIndex.removeExpiredLabels(rule, now) {
			continue
		}
		if len(rule.Labels) == 0 {
			err = l.DeleteLabelRuleLocked(key)
		} else {
			err = l.SaveLabelRuleLocked(rule)
		}
		if err != nil {
			log.Error("failed to save rule expired label rule", zap.String("rule-key", key), zap.Error(err))
		}
	}
}

func (l *RegionLabeler) mustLoadRules() bool {
	startKey := ""
	// pending accumulates loaded rules across pages and retries. It is only
	// touched by this goroutine, so it does not need l's lock until the
	// final publish below.
	pending := make(map[string]*LabelRule)
	wait := 10 * time.Second
	failpoint.Inject("fastRetryLoad", func() {
		wait = 100 * time.Millisecond
	})
	start := time.Now()
	for {
		err := l.loadAllRules(&startKey, pending)
		if err == nil {
			break
		}
		log.Info("load rules failed, retrying", zap.String("start-key", startKey), zap.Error(err))
		select {
		case <-l.ctx.Done():
			log.Warn("exit load rules due to context cancel")
			return false
		case <-time.After(wait):
		}
		if time.Since(start) > 10*time.Minute {
			log.Warn("load rules too slow", zap.Duration("cost", time.Since(start)))
		}
	}

	// This is a best-effort early exit to skip a pointless publish; it is not
	// what makes cancellation safe. l.ctx (or its parent) can be canceled at
	// any point without ever touching l's lock, so the actual guarantee that
	// a canceled labeler never reports ready lives in IsReady() itself, which
	// re-checks the context on every call.
	if l.ctx.Err() != nil {
		log.Warn("exit load rules due to context cancel")
		return false
	}
	// Build every in-memory index against pending while it is still local to
	// this goroutine. The lock then covers only the final field swap instead
	// of an O(len(pending)) rebuild, and readers cannot observe mismatched
	// rule, range, or keyspace views.
	index := newLabelRuleIndex(pending)
	l.Lock()
	l.ruleIndex = index
	l.Unlock()
	log.Info("new region labeler created", zap.Int("label-rules-count", len(pending)), zap.Duration("cost", time.Since(start)))
	// Test-only hooks: enteredPauseBeforeReady lets a test know execution has
	// actually reached this point (instead of guessing with a sleep) before
	// racing a cancellation in via pauseBeforeReady, right after the indexes
	// have been published but before isReady is actually stored. IsReady()'s
	// own re-check of l.ctx (not this call site) is what makes that race safe.
	failpoint.InjectCall("enteredPauseBeforeReady")
	failpoint.Inject("pauseBeforeReady", func() {
		time.Sleep(300 * time.Millisecond)
	})
	l.isReady.Store(true)
	return true
}

// loadAllRules loads all region label rules into pending, resuming from
// *startKey on retry. The rule indexes are published only once the whole scan
// succeeds (see mustLoadRules), so a partial or failed load never becomes
// visible outside this goroutine.
func (l *RegionLabeler) loadAllRules(startKey *string, pending map[string]*LabelRule) error {
	var toDelete []string
	var rules []*LabelRule
	loadedCount := 0
	batch := 100
	for {
		// Check between pages so a cancellation during a long, otherwise
		// successful multi-page scan (millions of rules, thousands of pages)
		// is noticed within one page instead of only after the whole scan
		// finishes; Close()/Stop() would otherwise have to wait for every
		// remaining page to be read before they can return.
		if err := l.ctx.Err(); err != nil {
			return err
		}
		failpoint.Inject("slowLoadPage", func() {
			time.Sleep(100 * time.Millisecond)
		})
		err := l.storage.LoadRegionRulesWithStartKey(*startKey, batch, func(k, v string) {
			*startKey = k
			loadedCount++
			r, err := NewLabelRuleFromJSON([]byte(v))
			if err != nil {
				log.Warn("failed to unmarshal label rule value", zap.String("rule-key", k), zap.String("rule-value", v), errs.ZapError(errs.ErrLoadRule))
				toDelete = append(toDelete, k)
				return
			}
			err = r.checkAndAdjust()
			if err != nil {
				log.Warn("failed to adjust label rule", zap.String("rule-key", k), zap.String("rule-value", v), zap.Error(err))
				toDelete = append(toDelete, k)
				return
			}
			rules = append(rules, r)
		})

		for _, id := range toDelete {
			if err := l.storage.RunInTxn(l.ctx, func(txn kv.Txn) error {
				return l.storage.DeleteRegionRule(txn, id)
			}); err != nil {
				if errors.Is(err, context.Canceled) {
					return err
				}
				log.Warn("failed to delete region rule", zap.String("rule-key", id), zap.Error(err))
			}
		}

		for _, rule := range rules {
			pending[rule.ID] = rule
		}

		if err != nil {
			return err
		}
		// Test-only hook: lets tests wait for a page to actually finish
		// before racing a cancellation in, instead of guessing with a sleep.
		// A no-op outside tests (EnableCall is never called in production).
		failpoint.InjectCall("pageLoaded")

		if loadedCount < batch {
			return nil
		}
		loadedCount = 0
		rules = rules[:0]
		toDelete = toDelete[:0]
	}
}

func (l *RegionLabeler) loadRules() error {
	var toDelete []string
	loadedRules := make(map[string]*LabelRule)
	start := time.Now()
	err := l.storage.LoadRegionRules(func(k, v string) {
		r, err := NewLabelRuleFromJSON([]byte(v))
		if err != nil {
			log.Warn("failed to unmarshal label rule value", zap.String("rule-key", k), zap.String("rule-value", v), errs.ZapError(errs.ErrLoadRule))
			toDelete = append(toDelete, k)
			return
		}
		err = r.checkAndAdjust()
		if err != nil {
			log.Warn("failed to adjust label rule", zap.String("rule-key", k), zap.String("rule-value", v), zap.Error(err))
			toDelete = append(toDelete, k)
			return
		}
		loadedRules[r.ID] = r
	})
	if err != nil {
		return err
	}
	for _, id := range toDelete {
		if err := l.storage.RunInTxn(l.ctx, func(txn kv.Txn) error {
			return l.storage.DeleteRegionRule(txn, id)
		}); err != nil {
			return err
		}
	}
	index := newLabelRuleIndex(loadedRules)
	l.Lock()
	l.ruleIndex = index
	l.Unlock()
	log.Info("load region rules finished", zap.Float64("take", time.Since(start).Seconds()))
	return nil
}

// BuildRangeListLocked publishes pending derived rule-index updates.
// Deprecated: Unlock publishes them automatically.
func (l *RegionLabeler) BuildRangeListLocked() {
	l.ruleIndex.buildRanges()
}

// GetSplitKeys returns all split keys in the range (start, end).
func (l *RegionLabeler) GetSplitKeys(start, end []byte) [][]byte {
	if !l.IsReady() {
		return nil
	}
	l.RLock()
	defer l.RUnlock()
	return l.ruleIndex.getSplitKeys(start, end)
}

func filterExpiredLabels(rule *LabelRule, now time.Time) *LabelRule {
	if rule == nil {
		return nil
	}

	hasExpired := false
	for i := range rule.Labels {
		if rule.Labels[i].expireBefore(now) {
			hasExpired = true
			break
		}
	}
	if !hasExpired {
		return rule
	}

	labels := make([]RegionLabel, 0, len(rule.Labels))
	for i := range rule.Labels {
		if !rule.Labels[i].expireBefore(now) {
			labels = append(labels, rule.Labels[i])
		}
	}
	if len(labels) == 0 {
		return nil
	}
	return &LabelRule{
		ID:       rule.ID,
		Index:    rule.Index,
		RuleType: rule.RuleType,
		Data:     rule.Data,
		Labels:   labels,
	}
}

// GetAllLabelRules returns all the rules.
func (l *RegionLabeler) GetAllLabelRules() ([]*LabelRule, error) {
	if !l.IsReady() {
		return nil, errs.ErrRegionLabelerNotReady
	}
	l.RLock()
	defer l.RUnlock()

	now := time.Now()
	rules := make([]*LabelRule, 0, len(l.ruleIndex.rules))
	for _, rule := range l.ruleIndex.rules {
		if filteredRule := filterExpiredLabels(rule, now); filteredRule != nil {
			rules = append(rules, filteredRule)
		}
	}
	return rules, nil
}

// GetLabelRules returns the rules that match the given ids.
func (l *RegionLabeler) GetLabelRules(ids []string) ([]*LabelRule, error) {
	if !l.IsReady() {
		return nil, errs.ErrRegionLabelerNotReady
	}
	l.RLock()
	defer l.RUnlock()

	now := time.Now()
	rules := make([]*LabelRule, 0, len(ids))
	for _, id := range ids {
		if rule, ok := l.ruleIndex.rules[id]; ok {
			if filteredRule := filterExpiredLabels(rule, now); filteredRule != nil {
				rules = append(rules, filteredRule)
			}
		}
	}
	return rules, nil
}

// GetLabelRule returns the Rule with the same ID. A nil rule with a nil
// error means the labeler is ready but no such rule exists; a nil rule with
// a non-nil error means the labeler is not ready yet, so callers should not
// treat the two cases as equivalent (e.g. HTTP handlers should not map both
// to 404).
func (l *RegionLabeler) GetLabelRule(id string) (*LabelRule, error) {
	if !l.IsReady() {
		return nil, errs.ErrRegionLabelerNotReady
	}
	l.RLock()
	defer l.RUnlock()
	return filterExpiredLabels(l.ruleIndex.rules[id], time.Now()), nil
}

// len only used for test.
func (l *RegionLabeler) len() int {
	l.Lock()
	defer l.Unlock()
	return len(l.ruleIndex.rules)
}

// SetLabelRule inserts or updates a LabelRule.
func (l *RegionLabeler) SetLabelRule(rule *LabelRule) error {
	l.Lock()
	defer l.Unlock()
	return l.SetLabelRuleLocked(rule)
}

// SetLabelRuleLocked inserts or updates a LabelRule. The enclosing Unlock
// publishes derived index updates.
func (l *RegionLabeler) SetLabelRuleLocked(rule *LabelRule) error {
	if !l.IsReady() {
		return errs.ErrRegionLabelerNotReady
	}
	if err := rule.checkAndAdjust(); err != nil {
		return err
	}
	if err := l.SaveLabelRuleLocked(rule); err != nil {
		return err
	}
	l.ruleIndex.set(rule)
	return nil
}

// SaveLabelRuleLocked inserts or updates a LabelRule but not buildRangeList.
// It only saves the rule to storage, and does not update the in-memory states.
func (l *RegionLabeler) SaveLabelRuleLocked(rule *LabelRule) error {
	if !l.IsReady() {
		return errs.ErrRegionLabelerNotReady
	}
	return l.storage.RunInTxn(l.ctx, func(txn kv.Txn) error {
		return l.storage.SaveRegionRule(txn, rule.ID, rule)
	})
}

// DeleteLabelRule removes a LabelRule.
func (l *RegionLabeler) DeleteLabelRule(id string) error {
	if !l.IsReady() {
		return errs.ErrRegionLabelerNotReady
	}
	l.Lock()
	defer l.Unlock()
	if _, ok := l.ruleIndex.rules[id]; !ok {
		return errs.ErrRegionRuleNotFound.FastGenByArgs(id)
	}
	return l.DeleteLabelRuleLocked(id)
}

// DeleteLabelRuleLocked removes a LabelRule. The enclosing Unlock publishes
// derived index updates.
func (l *RegionLabeler) DeleteLabelRuleLocked(id string) error {
	if !l.IsReady() {
		return errs.ErrRegionLabelerNotReady
	}
	if err := l.storage.RunInTxn(l.ctx, func(txn kv.Txn) error {
		return l.storage.DeleteRegionRule(txn, id)
	}); err != nil {
		return err
	}
	l.ruleIndex.delete(id)
	return nil
}

// Patch updates multiple region rules in a batch.
func (l *RegionLabeler) Patch(patch LabelRulePatch) error {
	if !l.IsReady() {
		return errs.ErrRegionLabelerNotReady
	}
	// setRulesMap is used to solve duplicate entries in DeleteRules and SetRules.
	// Note: We maintain compatibility with the previous behavior, which is to process DeleteRules before SetRules
	// If there are duplicate rules, we will prioritize SetRules and select the last one from SetRules.
	setRulesMap := make(map[string]*LabelRule)

	for _, rule := range patch.SetRules {
		if err := rule.checkAndAdjust(); err != nil {
			return err
		}
		setRulesMap[rule.ID] = rule
	}

	// save to storage
	var batch []func(kv.Txn) error
	for _, key := range patch.DeleteRules {
		if _, ok := setRulesMap[key]; ok {
			continue
		}
		localKey := key
		batch = append(batch, func(txn kv.Txn) error {
			return l.storage.DeleteRegionRule(txn, localKey)
		})
	}
	for _, rule := range setRulesMap {
		localID, localRule := rule.ID, rule
		batch = append(batch, func(txn kv.Txn) error {
			return l.storage.SaveRegionRule(txn, localID, localRule)
		})
	}
	if err := endpoint.RunBatchOpInTxn(l.ctx, l.storage, batch); err != nil {
		return err
	}

	// update in-memory states.
	l.Lock()
	defer l.Unlock()

	for _, key := range patch.DeleteRules {
		l.ruleIndex.delete(key)
	}
	for _, rule := range setRulesMap {
		l.ruleIndex.set(rule)
	}
	return nil
}

// GetRegionLabel returns the label of the region for a key.
// If there are multiple rules that match the key, the one with max rule index will be returned.
func (l *RegionLabeler) GetRegionLabel(region *core.RegionInfo, key string) string {
	if !l.IsReady() {
		return ""
	}
	l.RLock()
	defer l.RUnlock()
	now := time.Now()
	value, index := "", -1
	// search ranges
	rules, keyspaceRule, ok := l.ruleIndex.getRangeRules(region.GetStartKey(), region.GetEndKey())
	if !ok {
		return ""
	}
	applyRule := func(r *LabelRule) {
		if r.Index <= index && value != "" {
			return
		}
		for _, label := range r.Labels {
			if label.expireBefore(now) {
				continue
			}
			if label.Key == key {
				value, index = label.Value, r.Index
			}
		}
	}
	for _, rule := range rules {
		applyRule(rule.(*LabelRule))
	}
	if keyspaceRule != nil {
		applyRule(keyspaceRule)
	}
	return value
}

// ScheduleDisabled returns true if the region is lablelld with schedule-disabled.
func (l *RegionLabeler) ScheduleDisabled(region *core.RegionInfo) bool {
	if !l.IsReady() {
		return true
	}
	v := l.GetRegionLabel(region, scheduleOptionLabel)
	return strings.EqualFold(v, scheduleOptionValueDeny)
}

// GetRegionLabels returns the labels of the region.
// For each key, the label with max rule index will be returned.
func (l *RegionLabeler) GetRegionLabels(region *core.RegionInfo) []*RegionLabel {
	if !l.IsReady() {
		return nil
	}
	l.RLock()
	defer l.RUnlock()
	type valueIndex struct {
		value string
		index int
	}
	labels := make(map[string]valueIndex)
	now := time.Now()
	// search ranges
	rules, keyspaceRule, ok := l.ruleIndex.getRangeRules(region.GetStartKey(), region.GetEndKey())
	if ok {
		applyRule := func(r *LabelRule) {
			for _, label := range r.Labels {
				if label.expireBefore(now) {
					continue
				}
				if old, ok := labels[label.Key]; !ok || old.index < r.Index {
					labels[label.Key] = valueIndex{label.Value, r.Index}
				}
			}
		}
		for _, rule := range rules {
			applyRule(rule.(*LabelRule))
		}
		if keyspaceRule != nil {
			applyRule(keyspaceRule)
		}
	}
	result := make([]*RegionLabel, 0, len(labels))
	for k, l := range labels {
		result = append(result, &RegionLabel{
			Key:   k,
			Value: l.value,
		})
	}
	return result
}

// Close stops the background rule loading and waits for it to exit.
func (l *RegionLabeler) Close() {
	l.cancel()
	l.wg.Wait()
}

// MakeKeyRanges is a helper function to make key ranges.
func MakeKeyRanges(keys ...string) []any {
	var res []any
	for i := 0; i < len(keys); i += 2 {
		res = append(res, map[string]any{"start_key": keys[i], "end_key": keys[i+1]})
	}
	return res
}