// Copyright 2017 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package schedulers

import (
	"bytes"
	"strconv"
	"time"

	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/log"
	"github.com/tikv/pd/pkg/errs"
	"github.com/tikv/pd/server/core"
	"github.com/tikv/pd/server/schedule"
	"github.com/tikv/pd/server/schedule/filter"
	"github.com/tikv/pd/server/schedule/operator"
	"github.com/tikv/pd/server/schedule/opt"
	"github.com/tikv/pd/server/schedule/selector"
	"go.uber.org/zap"
)

const (
	scanLimit                    = 1000
	defaultAdjacentPeerLimit     = 1
	defaultAdjacentLeaderLimit   = 64
	minAdjacentSchedulerInterval = time.Second
	maxAdjacentSchedulerInterval = 30 * time.Second
	// AdjacentRegionName is balance adjacent region scheduler name.
	AdjacentRegionName = "balance-adjacent-region-scheduler"
	// AdjacentRegionType is balance adjacent region scheduler type.
	AdjacentRegionType = "adjacent-region"
)

func init() {
	schedule.RegisterSliceDecoderBuilder(AdjacentRegionType, func(args []string) schedule.ConfigDecoder {
		return func(v interface{}) error {
			conf, ok := v.(*balanceAdjacentRegionConfig)
			if !ok {
				return errs.ErrScheduleConfigNotExist.FastGenByArgs()
			}
			if len(args) == 2 {
				leaderLimit, err := strconv.ParseUint(args[0], 10, 64)
				if err != nil {
					return errs.ErrStrconvParseUint.Wrap(err).FastGenWithCause()
				}
				peerLimit, err := strconv.ParseUint(args[1], 10, 64)
				if err != nil {
					return errs.ErrStrconvParseUint.Wrap(err).FastGenWithCause()
				}
				conf.LeaderLimit = leaderLimit
				conf.PeerLimit = peerLimit
				return nil
			}
			conf.LeaderLimit = defaultAdjacentLeaderLimit
			conf.PeerLimit = defaultAdjacentPeerLimit
			conf.Name = AdjacentRegionName
			return nil
		}
	})

	schedule.RegisterScheduler(AdjacentRegionType, func(opController *schedule.OperatorController, storage *core.Storage, decoder schedule.ConfigDecoder) (schedule.Scheduler, error) {
		conf := &balanceAdjacentRegionConfig{
			LeaderLimit: defaultAdjacentLeaderLimit,
			PeerLimit:   defaultAdjacentPeerLimit,
		}
		if err := decoder(conf); err != nil {
			return nil, err
		}
		return newBalanceAdjacentRegionScheduler(opController, conf), nil
	})
}

type balanceAdjacentRegionConfig struct {
	Name        string `json:"name"`
	LeaderLimit uint64 `json:"leader-limit"`
	PeerLimit   uint64 `json:"peer-limit"`
}

// balanceAdjacentRegionScheduler will disperse adjacent regions.
// we will scan a part regions order by key, then select the longest
// adjacent regions and disperse them. finally, we will guarantee
// 1. any two adjacent regions' leader will not in the same store
// 2. the two regions' leader will not in the public store of this two regions
type balanceAdjacentRegionScheduler struct {
	*BaseScheduler
	selector             *selector.RandomSelector
	lastKey              []byte
	cacheRegions         *adjacentState
	conf                 *balanceAdjacentRegionConfig
	adjacentRegionsCount int
}

type adjacentState struct {
	assignedStoreIds []uint64
	regions          []*core.RegionInfo
	head             int
}

func (a *adjacentState) clear() {
	a.assignedStoreIds = a.assignedStoreIds[:0]
	a.regions = a.regions[:0]
	a.head = 0
}

func (a *adjacentState) len() int {
	return len(a.regions) - a.head
}

// newBalanceAdjacentRegionScheduler creates a scheduler that tends to disperse adjacent region
// on each store.
func newBalanceAdjacentRegionScheduler(opController *schedule.OperatorController, conf *balanceAdjacentRegionConfig) schedule.Scheduler {
	filters := []filter.Filter{
		&filter.StoreStateFilter{ActionScope: AdjacentRegionName, TransferLeader: true, MoveRegion: true},
		filter.NewSpecialUseFilter(AdjacentRegionName),
	}
	base := NewBaseScheduler(opController)
	s := &balanceAdjacentRegionScheduler{
		BaseScheduler: base,
		selector:      selector.NewRandomSelector(filters),
		conf:          conf,
		lastKey:       []byte(""),
	}
	return s
}

func (l *balanceAdjacentRegionScheduler) GetName() string {
	return l.conf.Name
}

func (l *balanceAdjacentRegionScheduler) GetType() string {
	return AdjacentRegionType
}

func (l *balanceAdjacentRegionScheduler) EncodeConfig() ([]byte, error) {
	return schedule.EncodeConfig(l.conf)
}

func (l *balanceAdjacentRegionScheduler) GetMinInterval() time.Duration {
	return minAdjacentSchedulerInterval
}

func (l *balanceAdjacentRegionScheduler) GetNextInterval(interval time.Duration) time.Duration {
	return intervalGrow(interval, maxAdjacentSchedulerInterval, linearGrowth)
}

func (l *balanceAdjacentRegionScheduler) IsScheduleAllowed(cluster opt.Cluster) bool {
	return l.allowBalanceLeader() || l.allowBalancePeer()
}

func (l *balanceAdjacentRegionScheduler) allowBalanceLeader() bool {
	return l.OpController.OperatorCount(operator.OpAdjacent|operator.OpLeader) < l.conf.LeaderLimit
}

func (l *balanceAdjacentRegionScheduler) allowBalancePeer() bool {
	return l.OpController.OperatorCount(operator.OpAdjacent|operator.OpRegion) < l.conf.PeerLimit
}

func (l *balanceAdjacentRegionScheduler) Schedule(cluster opt.Cluster) []*operator.Operator {
	if l.cacheRegions == nil {
		l.cacheRegions = &adjacentState{
			assignedStoreIds: make([]uint64, 0, len(cluster.GetStores())),
			regions:          make([]*core.RegionInfo, 0, scanLimit),
			head:             0,
		}
	}
	// we will process cache firstly
	if l.cacheRegions.len() >= 2 {
		return l.process(cluster)
	}

	l.cacheRegions.clear()
	regions := cluster.ScanRegions(l.lastKey, nil, scanLimit)
	// scan to the end
	if len(regions) <= 1 {
		schedulerStatus.WithLabelValues(l.GetName(), "adjacent-count").Set(float64(l.adjacentRegionsCount))
		l.adjacentRegionsCount = 0
		l.lastKey = []byte("")
		return nil
	}

	// calculate max adjacentRegions and record to the cache
	adjacentRegions := make([]*core.RegionInfo, 0, scanLimit)
	adjacentRegions = append(adjacentRegions, regions[0])
	maxLen := 0
	for i, r := range regions[1:] {
		l.lastKey = r.GetStartKey()

		// append if the region are adjacent
		lastRegion := adjacentRegions[len(adjacentRegions)-1]
		if lastRegion.GetLeader().GetStoreId() == r.GetLeader().GetStoreId() && bytes.Equal(lastRegion.GetEndKey(), r.GetStartKey()) {
			adjacentRegions = append(adjacentRegions, r)
			if i != len(regions)-2 { // not the last element
				continue
			}
		}

		if len(adjacentRegions) == 1 {
			adjacentRegions[0] = r
		} else {
			// got an max length adjacent regions in this range
			if maxLen < len(adjacentRegions) {
				l.cacheRegions.clear()
				maxLen = len(adjacentRegions)
				l.cacheRegions.regions = append(l.cacheRegions.regions, adjacentRegions...)
				adjacentRegions = adjacentRegions[:0]
				adjacentRegions = append(adjacentRegions, r)
			}
		}
	}

	l.adjacentRegionsCount += maxLen
	return l.process(cluster)
}

func (l *balanceAdjacentRegionScheduler) process(cluster opt.Cluster) []*operator.Operator {
	if l.cacheRegions.len() < 2 {
		return nil
	}
	head := l.cacheRegions.head
	r1 := l.cacheRegions.regions[head]
	r2 := l.cacheRegions.regions[head+1]

	defer func() {
		if l.cacheRegions.len() < 0 {
			log.Fatal("cache overflow", zap.String("scheduler", l.GetName()), errs.ZapError(errs.ErrCacheOverflow))
		}
		l.cacheRegions.head = head + 1
		l.lastKey = r2.GetStartKey()
	}()
	// after the cluster is prepared, there is a gap that some regions heartbeats are not received.
	// Leader of those region is nil, and we should skip them.
	if r1.GetLeader() == nil || r2.GetLeader() == nil || l.unsafeToBalance(cluster, r1) {
		schedulerCounter.WithLabelValues(l.GetName(), "skip").Inc()
		return nil
	}
	op := l.disperseLeader(cluster, r1, r2)
	if op == nil {
		schedulerCounter.WithLabelValues(l.GetName(), "no-leader").Inc()
		op = l.dispersePeer(cluster, r1)
	}
	if op == nil {
		schedulerCounter.WithLabelValues(l.GetName(), "no-peer").Inc()
		l.cacheRegions.assignedStoreIds = l.cacheRegions.assignedStoreIds[:0]
		return nil
	}
	op.Counters = append(op.Counters, schedulerCounter.WithLabelValues(l.GetName(), "new-operator"))
	return []*operator.Operator{op}
}

func (l *balanceAdjacentRegionScheduler) unsafeToBalance(cluster opt.Cluster, region *core.RegionInfo) bool {
	if !opt.IsRegionReplicated(cluster, region) {
		return true
	}
	leaderStoreID := region.GetLeader().GetStoreId()
	store := cluster.GetStore(leaderStoreID)
	if store == nil {
		log.Error("failed to get the store", zap.Uint64("store-id", leaderStoreID), errs.ZapError(errs.ErrGetSourceStore))
		return true
	}
	s := l.selector.SelectSource(cluster, []*core.StoreInfo{store})
	if s == nil {
		return true
	}
	// Skip hot regions.
	if cluster.IsRegionHot(region) {
		schedulerCounter.WithLabelValues(l.GetName(), "region-hot").Inc()
		return true
	}
	return false
}

func (l *balanceAdjacentRegionScheduler) disperseLeader(cluster opt.Cluster, before *core.RegionInfo, after *core.RegionInfo) *operator.Operator {
	if !l.allowBalanceLeader() {
		return nil
	}
	diffPeers := before.GetDiffFollowers(after)
	if len(diffPeers) == 0 {
		return nil
	}
	storesInfo := make([]*core.StoreInfo, 0, len(diffPeers))
	for _, p := range diffPeers {
		if store := cluster.GetStore(p.GetStoreId()); store != nil {
			storesInfo = append(storesInfo, store)
		}
	}
	target := l.selector.SelectTarget(cluster, storesInfo)
	if target == nil {
		return nil
	}
	op, err := operator.CreateTransferLeaderOperator("balance-adjacent-leader", cluster, before, before.GetLeader().GetStoreId(), target.GetID(), operator.OpAdjacent)
	if err != nil {
		log.Debug("fail to create transfer leader operator", errs.ZapError(err))
		return nil
	}
	op.SetPriorityLevel(core.LowPriority)
	schedulerCounter.WithLabelValues(l.GetName(), "adjacent-leader").Inc()
	return op
}

func (l *balanceAdjacentRegionScheduler) dispersePeer(cluster opt.Cluster, region *core.RegionInfo) *operator.Operator {
	if !l.allowBalancePeer() {
		return nil
	}
	// scoreGuard guarantees that the distinct score will not decrease.
	leaderStoreID := region.GetLeader().GetStoreId()
	stores := cluster.GetRegionStores(region)
	source := cluster.GetStore(leaderStoreID)
	if source == nil {
		log.Error("failed to get the source store", zap.Uint64("store-id", leaderStoreID), errs.ZapError(errs.ErrGetSourceStore))
		return nil
	}
	var scoreGuard filter.Filter
	if cluster.IsPlacementRulesEnabled() {
		scoreGuard = filter.NewRuleFitFilter(l.GetName(), cluster, region, leaderStoreID)
	} else {
		scoreGuard = filter.NewDistinctScoreFilter(l.GetName(), cluster.GetLocationLabels(), stores, source)
	}
	excludeStores := region.GetStoreIds()
	for _, storeID := range l.cacheRegions.assignedStoreIds {
		if _, ok := excludeStores[storeID]; !ok {
			excludeStores[storeID] = struct{}{}
		}
	}

	filters := []filter.Filter{
		filter.NewExcludedFilter(l.GetName(), nil, excludeStores),
		scoreGuard,
	}
	target := l.selector.SelectTarget(cluster, cluster.GetStores(), filters...)
	if target == nil {
		return nil
	}
	newPeer := &metapb.Peer{StoreId: target.GetID()}
	// record the store id and exclude it in next time
	l.cacheRegions.assignedStoreIds = append(l.cacheRegions.assignedStoreIds, newPeer.GetStoreId())

	op, err := operator.CreateMovePeerOperator("balance-adjacent-peer", cluster, region, operator.OpAdjacent, leaderStoreID, newPeer)
	if err != nil {
		schedulerCounter.WithLabelValues(l.GetName(), "create-operator-fail").Inc()
		return nil
	}
	op.SetPriorityLevel(core.LowPriority)
	schedulerCounter.WithLabelValues(l.GetName(), "adjacent-peer").Inc()
	return op
}