package round_robin

import (
	"google.golang.org/grpc/balancer"
	"google.golang.org/grpc/balancer/base"
	"math"
	"sync"
)

/*
加权轮询
*/
type WeightBalancerBuilder struct {
}

func (w *WeightBalancerBuilder) Build(info base.PickerBuildInfo) balancer.Picker {
	nodes := make([]*weightNode, 0, len(info.ReadySCs))

	for k, v := range info.ReadySCs {
		weight, ok := v.Address.Attributes.Value("weight").(uint32)
		if !ok {
			panic("weight must be set")
		}

		nodes = append(nodes, &weightNode{
			weight:        weight,
			currentWeight: weight,
			effectWeight:  weight,

			c: k,
		})
	}

	p := &weightBalancer{nodes: nodes}
	return p
}

type weightBalancer struct {
	nodes []*weightNode
}

func (w *weightBalancer) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
	if len(w.nodes) == 0 {
		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
	}

	var (
		totalWeight uint32
		res         *weightNode
	)

	for _, e := range w.nodes {
		e.Lock()
		totalWeight += e.effectWeight
		e.currentWeight = e.currentWeight + e.effectWeight
		if res == nil || res.currentWeight < e.currentWeight {
			res = e
		}
		e.Unlock()
	}

	res.Lock()
	res.currentWeight = res.currentWeight - totalWeight
	res.Unlock()

	return balancer.PickResult{
		SubConn: res.c,
		Done: func(info balancer.DoneInfo) {
			res.Lock()
			if info.Err != nil && res.effectWeight == 0 {
				return
			}

			if info.Err == nil && res.effectWeight == math.MaxUint32 {
				return
			}

			if info.Err != nil {
				res.effectWeight--
			} else {
				res.effectWeight++
			}
			res.Unlock()
		},
	}, nil
}

type weightNode struct {
	weight        uint32
	currentWeight uint32
	effectWeight  uint32

	c balancer.SubConn

	sync.Mutex
}