package slprober

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func Test_probeNextGenSvc(t *testing.T) {
	type args struct {
		clusterID       string
		groupName       string
		groupType       string
		svc             *v1.Service
		handshakeErrors map[string]error
		expectedSuccess bool
	}
	tests := []struct {
		name string
		args args
	}{
		{
			name: "tidbgroup, alive",
			args: args{
				clusterID: "123",
				groupType: "tidbGroup",
				groupName: "db-456",
				svc: &v1.Service{
					ObjectMeta: metav1.ObjectMeta{
						Namespace: "tidb123",
						Name:      "db-456-tidb",
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				handshakeErrors: map[string]error{
					"1.2.3.4:4000": nil,
				},
				expectedSuccess: true,
			},
		},
		{
			name: "tidbgroup, dead",
			args: args{
				clusterID: "123",
				groupType: "tidbGroup",
				groupName: "db-456",
				svc: &v1.Service{
					ObjectMeta: metav1.ObjectMeta{
						Namespace: "tidb123",
						Name:      "db-456-tidb",
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				handshakeErrors: map[string]error{
					"1.2.3.4:4000": errors.New("connection failed"),
				},
				expectedSuccess: false,
			},
		},
		{
			name: "tidbgroup, timeout",
			args: args{
				clusterID: "123",
				groupType: "tidbGroup",
				groupName: "db-456",
				svc: &v1.Service{
					ObjectMeta: metav1.ObjectMeta{
						Namespace: "tidb123",
						Name:      "db-456-tidb",
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				handshakeErrors: map[string]error{
					"1.2.3.4:4000": context.DeadlineExceeded,
				},
				expectedSuccess: false,
			},
		},
		{
			name: "tiproxygroup, alive",
			args: args{
				clusterID: "123",
				groupType: "tiproxyGroup",
				groupName: "db-789",
				svc: &v1.Service{
					ObjectMeta: metav1.ObjectMeta{
						Namespace: "tidb123",
						Name:      "db-789-tiproxy",
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "5.6.7.8",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				handshakeErrors: map[string]error{
					"5.6.7.8:4000": nil,
				},
				expectedSuccess: true,
			},
		},
		{
			name: "tiproxygroup, dead",
			args: args{
				clusterID: "123",
				groupType: "tiproxyGroup",
				groupName: "db-789",
				svc: &v1.Service{
					ObjectMeta: metav1.ObjectMeta{
						Namespace: "tidb123",
						Name:      "db-789-tiproxy",
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "5.6.7.8",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				handshakeErrors: map[string]error{
					"5.6.7.8:4000": errors.New("connection failed"),
				},
				expectedSuccess: false,
			},
		},
		{
			name: "tiproxygroup, timeout",
			args: args{
				clusterID: "123",
				groupType: "tiproxyGroup",
				groupName: "db-789",
				svc: &v1.Service{
					ObjectMeta: metav1.ObjectMeta{
						Namespace: "tidb123",
						Name:      "db-789-tiproxy",
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "5.6.7.8",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				handshakeErrors: map[string]error{
					"5.6.7.8:4000": context.DeadlineExceeded,
				},
				expectedSuccess: false,
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mockHandshakeCli := new(MockHandshakeClient)
			shouldTimeout := false
			for addr, err := range tt.args.handshakeErrors {
				if errors.Is(err, context.DeadlineExceeded) {
					shouldTimeout = true
					mockHandshakeCli.On("TryConnect", mock.Anything, addr).After(1*time.Second).Return(nil, err)
				} else {
					mockHandshakeCli.On("TryConnect", mock.Anything, addr).Return(nil, err)
				}
			}
			resultCh := make(chan tidbGroupResult, 1)
			ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
			defer cancel()

			go probeNextGenSvc(
				ctx,
				tt.args.clusterID,
				tt.args.groupName,
				tt.args.groupType,
				tt.args.svc,
				resultCh,
				mockHandshakeCli,
			)
			select {
			case result := <-resultCh:
				if shouldTimeout {
					t.Fatal("Expected no result due to context cancellation")
				} else {
					assert.Equal(t, tt.args.expectedSuccess, result.success)
				}
			case <-ctx.Done():
				if !shouldTimeout {
					t.Fatal("test timed out")
				}
			}

			time.Sleep(200 * time.Millisecond) // Wait for goroutines to finish
			mockHandshakeCli.AssertExpectations(t)
		})
	}
}