// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
//go:build plus

package models

import (
	"github.com/TeaOSLab/EdgeCommon/pkg/dnsconfigs"
	"github.com/iwind/TeaGo/rands"
	"math/rand"
	"time"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

type recordIdInfo struct {
	Id     int64
	Weight int32
}

type RecordIds struct {
	IdList   []*recordIdInfo
	IdBucket []int64

	RoundIndex int

	totalWeight int64
}

func NewRecordIds() *RecordIds {
	return &RecordIds{}
}

func (this *RecordIds) IsEmpty() bool {
	return len(this.IdList) == 0
}

func (this *RecordIds) Add(newId int64, weight int32) {
	if weight <= 0 {
		weight = 10
	}

	const maxWeight = 999999
	if weight > maxWeight {
		weight = maxWeight
	}

	// 检查是否存在
	for _, idInfo := range this.IdList {
		if idInfo.Id == newId {
			return
		}
	}

	// 添加
	this.IdList = append(this.IdList, &recordIdInfo{
		Id:     newId,
		Weight: weight,
	})

	// 重置数据
	this.resetData()
}

func (this *RecordIds) Remove(oldId int64) {
	defer this.resetData()

	var newIdList = []*recordIdInfo{}
	for _, idInfo := range this.IdList {
		if idInfo.Id == oldId {
			continue
		}
		newIdList = append(newIdList, idInfo)
	}
	this.IdList = newIdList
}

// NextId for round-robin
func (this *RecordIds) NextId() int64 {
	var l = len(this.IdList)

	if l == 0 {
		return 0
	}

	if l == 1 {
		return this.IdList[0].Id
	}

	if this.RoundIndex > l-1 {
		this.RoundIndex = 0
	}

	var id = this.IdList[this.RoundIndex].Id

	this.RoundIndex++
	return id
}

func (this *RecordIds) RandomIds(count int) []int64 {
	if count <= 0 {
		count = dnsconfigs.NSAnswerDefaultSize
	}

	var totalRecords = len(this.IdList)

	if totalRecords == 0 {
		return nil
	}

	if totalRecords == 1 {
		return []int64{this.IdList[0].Id} // duplicate
	}

	if totalRecords < count {
		count = totalRecords
	}

	var totalIds = len(this.IdBucket)
	var startIndex = rands.Int(0, totalIds-1)
	var endIndex = startIndex + count - 1
	if endIndex <= totalIds-1 {
		return this.IdBucket[startIndex : endIndex+1]
	}
	return append(this.IdBucket[startIndex:totalIds], this.IdBucket[0:endIndex-totalIds+1]...)
}

func (this *RecordIds) resetData() {
	this.resetWeight()
}

func (this *RecordIds) resetWeight() {
	var totalWeight int64

	var weightMap = map[int32]bool{} // weight => bool
	var hasUniqueWeights = false
	var ids []int64

	for _, idInfo := range this.IdList {
		totalWeight += int64(idInfo.Weight)

		// 检查是否有不同的权重
		if len(weightMap) > 0 && !weightMap[idInfo.Weight] {
			hasUniqueWeights = true
		}
		weightMap[idInfo.Weight] = true

		ids = append(ids, idInfo.Id)
	}

	// 根据权重，重新组织IDs
	if hasUniqueWeights {
		var newIds = []int64{}
		for _, idInfo := range this.IdList {
			for i := int32(0); i < idInfo.Weight; i++ {
				newIds = append(newIds, idInfo.Id)
			}
		}
		ids = newIds
	}

	var countIds = len(ids)
	if countIds > 0 {
		rand.Shuffle(countIds, func(i, j int) {
			ids[i], ids[j] = ids[j], ids[i]
		})
	}

	this.totalWeight = totalWeight
	this.IdBucket = ids
}