package grpc

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

var no = 0

type WeightBalance struct {
	subConns   []balancer.SubConn
	subWeights []int64

	mu sync.Mutex
}

func (b *WeightBalance) Build(info base.PickerBuildInfo) balancer.Picker {

	if len(info.ReadySCs) <= 0 {
		return base.NewErrPicker(balancer.ErrNoSubConnAvailable)
	}

	conns := []balancer.SubConn{}
	weights := []int64{}
	i := 0
	for subConn, subInfo := range info.ReadySCs {
		i++
		conns = append(conns, subConn)

		//fmt.Println("subConn", subConn)
		//fmt.Println("Attributes", subInfo.Address.Attributes)
		//fmt.Println("Addr", subInfo.Address.Addr)

		weight := subInfo.Address.Attributes.Value("weight")
		w := int64(0)
		if weight != nil {
			w = weight.(int64)
		}

		if w <= 0 {
			w = 1
		}

		log.Printf("rpc discovery %d: %s ,%d", i, subInfo.Address.Addr, weight)
		weights = append(weights, w)
	}
	b.mu.Lock()
	b.subConns = conns
	b.subWeights = weights
	no = 0
	b.mu.Unlock()

	return b
}

func (b *WeightBalance) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
	b.mu.Lock()
	l := len(b.subWeights)
	no++
	no = no % l
	conn := b.subConns[no]
	b.mu.Unlock()

	return balancer.PickResult{
		SubConn: conn,
	}, nil
}

func (b *WeightBalance) Name() string {
	return "my_weight_round"
}