/*
 *
 * Copyright 2023 gRPC 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,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package test

import (
	"context"
	"errors"
	"fmt"
	"testing"

	"google.golang.org/grpc"
	"google.golang.org/grpc/balancer"
	"google.golang.org/grpc/balancer/base"
	"google.golang.org/grpc/connectivity"
	"google.golang.org/grpc/internal/balancer/stub"
	"google.golang.org/grpc/internal/stubserver"
	"google.golang.org/grpc/internal/testutils"
	testpb "google.golang.org/grpc/interop/grpc_testing"
	"google.golang.org/grpc/resolver"
)

type tsccPicker struct {
	sc balancer.SubConn
}

func (p *tsccPicker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
	return balancer.PickResult{SubConn: p.sc}, nil
}

// TestSubConnEmpty tests that removing all addresses from a SubConn and then
// re-adding them does not cause a panic and properly reconnects.
func (s) TestSubConnEmpty(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()

	// sc is the one SubConn used throughout the test.  Created on demand and
	// re-used on every update.
	var sc balancer.SubConn

	// Simple custom balancer that sets the address list to empty if the
	// resolver produces no addresses.  Pickfirst, by default, will remove the
	// SubConn in this case instead.
	bal := stub.BalancerFuncs{
		UpdateClientConnState: func(d *stub.BalancerData, ccs balancer.ClientConnState) error {
			if sc == nil {
				var err error
				sc, err = d.ClientConn.NewSubConn(ccs.ResolverState.Addresses, balancer.NewSubConnOptions{
					StateListener: func(state balancer.SubConnState) {
						switch state.ConnectivityState {
						case connectivity.Ready:
							d.ClientConn.UpdateState(balancer.State{
								ConnectivityState: connectivity.Ready,
								Picker:            &tsccPicker{sc: sc},
							})
						case connectivity.TransientFailure:
							d.ClientConn.UpdateState(balancer.State{
								ConnectivityState: connectivity.TransientFailure,
								Picker:            base.NewErrPicker(fmt.Errorf("error connecting: %v", state.ConnectionError)),
							})
						}
					},
				})
				if err != nil {
					t.Errorf("error creating initial subconn: %v", err)
				}
			} else {
				d.ClientConn.UpdateAddresses(sc, ccs.ResolverState.Addresses)
			}
			sc.Connect()

			if len(ccs.ResolverState.Addresses) == 0 {
				d.ClientConn.UpdateState(balancer.State{
					ConnectivityState: connectivity.TransientFailure,
					Picker:            base.NewErrPicker(errors.New("no addresses")),
				})
			} else {
				d.ClientConn.UpdateState(balancer.State{
					ConnectivityState: connectivity.Connecting,
					Picker:            &tsccPicker{sc: sc},
				})
			}
			return nil
		},
	}
	stub.Register("tscc", bal)

	// Start the stub server with our stub balancer.
	ss := &stubserver.StubServer{
		EmptyCallF: func(ctx context.Context, in *testpb.Empty) (*testpb.Empty, error) {
			return &testpb.Empty{}, nil
		},
	}
	if err := ss.Start(nil, grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"tscc":{}}]}`)); err != nil {
		t.Fatalf("Error starting server: %v", err)
	}
	defer ss.Stop()
	if _, err := ss.Client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
		t.Fatalf("EmptyCall failed: %v", err)
	}

	t.Log("Removing addresses from resolver and SubConn")
	ss.R.UpdateState(resolver.State{Addresses: []resolver.Address{}})
	testutils.AwaitState(ctx, t, ss.CC, connectivity.TransientFailure)

	t.Log("Re-adding addresses to resolver and SubConn")
	ss.R.UpdateState(resolver.State{Addresses: []resolver.Address{{Addr: ss.Address}}})
	if _, err := ss.Client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
		t.Fatalf("EmptyCall failed: %v", err)
	}
}