package grpclb

import (
	"context"
	"errors"
	"sync"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/balancer"
	grpclbstate "google.golang.org/grpc/balancer/grpclb/state"
	"google.golang.org/grpc/connectivity"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/grpclog"
	"google.golang.org/grpc/internal"
	"google.golang.org/grpc/internal/backoff"
	"google.golang.org/grpc/internal/resolver/dns"
	"google.golang.org/grpc/resolver"

	durationpb "github.com/golang/protobuf/ptypes/duration"
	lbpb "google.golang.org/grpc/balancer/grpclb/grpc_lb_v1"
)

const (
	lbTokenKey             = "lb-token"
	defaultFallbackTimeout = 10 * time.Second
	grpclbName             = "grpclb"
)

var errServerTerminatedConnection = errors.New("grpclb: failed to recv server list: server terminated connection")
var logger = grpclog.Component("grpclb")

func convertDuration(d *durationpb.Duration) time.Duration {
	if d == nil {
		return 0
	}
	return time.Duration(d.Seconds)*time.Second + time.Duration(d.Nanos)*time.Nanosecond
}

type loadBalancerClient struct {
	cc *grpc.ClientConn
}

func (c *loadBalancerClient) BalanceLoad(ctx context.Context, opts ...grpc.CallOption) (*balanceLoadClientStream, error) {
	desc := &grpc.StreamDesc{
		StreamName:    "BalanceLoad",
		ServerStreams: true,
		ClientStreams: true,
	}
	stream, err := c.cc.NewStream(ctx, desc, "/grpc.lb.v1.LoadBalancer/BalanceLoad", opts...)
	if err != nil {
		return nil, err
	}
	x := &balanceLoadClientStream{stream}
	return x, nil
}

type balanceLoadClientStream struct {
	grpc.ClientStream
}

func (x *balanceLoadClientStream) Send(m *lbpb.LoadBalanceRequest) error {
	return x.ClientStream.SendMsg(m)
}

func (x *balanceLoadClientStream) Recv() (*lbpb.LoadBalanceResponse, error) {
	m := new(lbpb.LoadBalanceResponse)
	if err := x.ClientStream.RecvMsg(m); err != nil {
		return nil, err
	}
	return m, nil
}

func init() {
	balancer.Register(newLBBuilder())
	dns.EnableSRVLookups = true
}

func newLBBuilder() balancer.Builder {
	return newLBBuilderWithFallbackTimeout(defaultFallbackTimeout)
}

func newLBBuilderWithFallbackTimeout(fallbackTimeout time.Duration) balancer.Builder {
	return &lbBuilder{
		fallbackTimeout: fallbackTimeout,
	}
}

type lbBuilder struct {
	fallbackTimeout time.Duration
}

func (b *lbBuilder) Name() string {
	return grpclbName
}

	r := &lbManualResolver{scheme: "grpclb-internal", ccb: cc}

	lb := &lbBalancer{
		cc:              newLBCacheClientConn(cc),
		target:          opt.Target.Endpoint,
		opt:             opt,
		fallbackTimeout: b.fallbackTimeout,
		doneCh:          make(chan struct{}),

		manualResolver: r,
		subConns:       make(map[resolver.Address]balancer.SubConn),
		scStates:       make(map[balancer.SubConn]connectivity.State),
		picker:         &errPicker{err: balancer.ErrNoSubConnAvailable},
		clientStats:    newRPCStats(),
		backoff:        backoff.DefaultExponential, // TODO: make backoff configurable.
	}

	var err error
	if opt.CredsBundle != nil {
		lb.grpclbClientConnCreds, err = opt.CredsBundle.NewWithMode(internal.CredsBundleModeBalancer)
		if err != nil {
			logger.Warningf("lbBalancer: client connection creds NewWithMode failed: %v", err)
		}
		lb.grpclbBackendCreds, err = opt.CredsBundle.NewWithMode(internal.CredsBundleModeBackendFromBalancer)
		if err != nil {
			logger.Warningf("lbBalancer: backend creds NewWithMode failed: %v", err)
		}
	}

	return lb
}

type lbBalancer struct {
	cc     *lbCacheClientConn
	target string
	opt    balancer.BuildOptions

	usePickFirst bool

	grpclbBackendCreds credentials.Bundle

	fallbackTimeout time.Duration
	doneCh          chan struct{}

	backoff backoff.Strategy

	clientStats *rpcStats

	state    connectivity.State
	subConns map[resolver.Address]balancer.SubConn   // Used to new/remove SubConn.
	scStates map[balancer.SubConn]connectivity.State // Used to filter READY SubConns.

	remoteBalancerConnected bool
	serverListReceived      bool

	resolvedBackendAddrs []resolver.Address
}

func (lb *lbBalancer) regeneratePicker(resetDrop bool) {
	if lb.state == connectivity.TransientFailure {
		lb.picker = &errPicker{err: balancer.ErrTransientFailure}
		return
	}

	if lb.state == connectivity.Connecting {
		lb.picker = &errPicker{err: balancer.ErrNoSubConnAvailable}
		return
	}

	var readySCs []balancer.SubConn
	if lb.usePickFirst {
		for _, sc := range lb.subConns {
			readySCs = append(readySCs, sc)
			break
		}
	} else {
		for _, a := range lb.backendAddrsWithoutMetadata {
			if sc, ok := lb.subConns[a]; ok {
				if st, ok := lb.scStates[sc]; ok && st == connectivity.Ready {
					readySCs = append(readySCs, sc)
				}
			}
		}
	}

		lb.picker = &errPicker{err: balancer.ErrNoSubConnAvailable}
		return
	}
	if lb.inFallback {
		lb.picker = newRRPicker(readySCs)
		return
	}
	if resetDrop {
		lb.picker = newLBPicker(lb.fullServerList, readySCs, lb.clientStats)
		return
	}
	prevLBPicker, ok := lb.picker.(*lbPicker)
	if !ok {
		lb.picker = newLBPicker(lb.fullServerList, readySCs, lb.clientStats)
		return
	}
	prevLBPicker.updateReadySCs(readySCs)
}

func (lb *lbBalancer) aggregateSubConnStates() connectivity.State {
	var numConnecting uint64

	for _, sc := range lb.subConns {
		if state, ok := lb.scStates[sc]; ok {
			switch state {
			case connectivity.Ready:
				return connectivity.Ready
			case connectivity.Connecting:
				numConnecting++
			}
		}
	}
	if numConnecting > 0 {
		return connectivity.Connecting
	}
	return connectivity.TransientFailure
}

func (lb *lbBalancer) UpdateSubConnState(sc balancer.SubConn, scs balancer.SubConnState) {
	s := scs.ConnectivityState
	if logger.V(2) {
		logger.Infof("lbBalancer: handle SubConn state change: %p, %v", sc, s)
	}
	lb.mu.Lock()
	defer lb.mu.Unlock()

	oldS, ok := lb.scStates[sc]
	if !ok {
		if logger.V(2) {
			logger.Infof("lbBalancer: got state changes for an unknown SubConn: %p, %v", sc, s)
		}
		return
	}
	lb.scStates[sc] = s
	switch s {
	case connectivity.Idle:
		sc.Connect()

		delete(lb.scStates, sc)

	lb.updateStateAndPicker((oldS == connectivity.Ready) != (s == connectivity.Ready), false)

	if lb.state != connectivity.Ready {

			lb.refreshSubConns(lb.resolvedBackendAddrs, true, lb.usePickFirst)
		}
	}
}

func (lb *lbBalancer) updateStateAndPicker(forceRegeneratePicker bool, resetDrop bool) {
	oldAggrState := lb.state

	if forceRegeneratePicker || (lb.state != oldAggrState) {
		lb.regeneratePicker(resetDrop)
	}

	lb.cc.UpdateState(balancer.State{ConnectivityState: lb.state, Picker: lb.picker})
}

func (lb *lbBalancer) fallbackToBackendsAfter(fallbackTimeout time.Duration) {
	timer := time.NewTimer(fallbackTimeout)
	defer timer.Stop()
	select {
	case <-timer.C:
	case <-lb.doneCh:
		return
	}
	lb.mu.Lock()
	if lb.inFallback || lb.serverListReceived {
		lb.mu.Unlock()
		return

	lb.refreshSubConns(lb.resolvedBackendAddrs, true, lb.usePickFirst)
	lb.mu.Unlock()
}

func (lb *lbBalancer) handleServiceConfig(gc *grpclbServiceConfig) {
	lb.mu.Lock()
	defer lb.mu.Unlock()

	newUsePickFirst := childIsPickFirst(gc)
	if lb.usePickFirst == newUsePickFirst {
		return
	}
	if logger.V(2) {
		logger.Infof("lbBalancer: switching mode, new usePickFirst: %+v", newUsePickFirst)
	}
	lb.refreshSubConns(lb.backendAddrs, lb.inFallback, newUsePickFirst)
}

}

func (lb *lbBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
	if logger.V(2) {
		logger.Infof("lbBalancer: UpdateClientConnState: %+v", ccs)
	}
	gc, _ := ccs.BalancerConfig.(*grpclbServiceConfig)
	lb.handleServiceConfig(gc)

	addrs := ccs.ResolverState.Addresses

	var remoteBalancerAddrs, backendAddrs []resolver.Address
	for _, a := range addrs {
		if a.Type == resolver.GRPCLB {
			a.Type = resolver.Backend
			remoteBalancerAddrs = append(remoteBalancerAddrs, a)
		} else {
			backendAddrs = append(backendAddrs, a)
		}
	}

		remoteBalancerAddrs = sd.BalancerAddresses
	}

		return balancer.ErrBadResolverState
	}

	if len(remoteBalancerAddrs) == 0 {
		if lb.ccRemoteLB != nil {
			lb.ccRemoteLB.close()
			lb.ccRemoteLB = nil
		}

		go lb.fallbackToBackendsAfter(lb.fallbackTimeout)
	}

		lb.manualResolver.UpdateState(resolver.State{Addresses: remoteBalancerAddrs})
	}

	lb.mu.Lock()
	lb.resolvedBackendAddrs = backendAddrs

		lb.refreshSubConns(lb.resolvedBackendAddrs, true, lb.usePickFirst)
	}
	lb.mu.Unlock()
	return nil
}

func (lb *lbBalancer) Close() {
	select {
	case <-lb.doneCh:
		return
	default:
	}
	close(lb.doneCh)
	if lb.ccRemoteLB != nil {
		lb.ccRemoteLB.close()
	}
	lb.cc.close()